Merge branch 'dev' into fix-lycoris-wrong-errors

CHANGELOG.md

@@ -1,26 +1,86 @@
 # Change Log for SD.Next
 
-## Update for 2025-06-06
+## Update for 2025-06-25
+
+- **Changes**  
+  - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support (in addition to existing JoyCaption Alpha)  
+  - Support Remote VAE with *Omnigen, Lumina 2 and PixArt*  
+  - Use Diffusers version of *OmniGen*  
+  - Control move global settings to control elements -> control settings tab  
+  - Control add setting to run hires with or without control  
+
+- **SDNQ Quantization**  
+  - Add modules_to_not_convert support for post mode  
+  - Fix Qwen 2.5 with int8 matmul  
+  - Fix Dora loading  
+  - Remove per layer GC  
+  - Improve offload compatibility  
+  - Add support for XYZ grid to test quantization modes  
+    *note*: you need to enable quantization and choose what it applies on, then xyz grid can change quantization mode  
+
+- **API**
+  - Add `/sdapi/v1/lora?lora=<lora_name>` endpoint that returns full lora info and metadata  
+  - Add `/sdapi/v1/controlnets?model_type=<model_type|all|None>` endpoints that returns list of available controlnets for specific model type  
+
+- **Fixes**  
+  - IPEX with DPM2++ FlowMatch samplers  
+  - Invalid attention processor with ControlNet  
+  - LTXVideo default scheduler  
+  - Balanced offload with OmniGen  
+  - Quantization with OmniGen  
+  - Do not save empty `params.txt` file  
+  - Override `params.txt` using `SD_PATH_PARAMS` env variable  
+  - Add `wheel` to requirements due to `pip` change  
+  - Case-insensitive sampler name matching  
+  - Fix delete file with gallery views  
+  - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen  
+  - Fix TAESD model type detection  
+
+## Update for 2025-06-16
+
+- **Feature**  
+  - Support for Python 3.13  
+  - TeaCache support for Lumina 2  
+  - Custom UNet and VAE loading support for Lumina 2  
 
 - **Changes**  
   - Increase the medvram mode threshold from 8GB to 12GB  
   - Set CPU backend to use FP32 by default  
+  - Relax Python version checks for Zluda  
+  - Make VAE options not require model reload  
+  - Add warning about incompatible attention processors  
 
 - **Torch**  
-  - set default to `torch==2.7.1`  
+  - Set default to `torch==2.7.1`  
+  - Force upgrade pip when installing Torch  
+
+- **ROCm**  
+  - Support ROCm 6.4 with `--use-nightly`  
+  - Don't override user set gfx version  
+  - Don't override gfx version with RX 9000  
+  - Fix flash-atten repo  
 
 - **SDNQ Quantization**  
   - Add group size support for convolutional layers  
   - Add quantized matmul support for for convolutional layers  
+  - Add 7-bit, 5-bit and 3-bit quantization support  
+  - Add separate quant mode option for Text Encoders  
   - Fix forced FP32 with tensorwise FP8 matmul  
   - Fix PyTorch <= 2.4 compatibility with FP8 matmul  
   - Fix VAE with conv quant  
+  - Don't ignore the Quantize with GPU option with offload mode `none` and `model`  
+  - High VRAM usage with Lumina 2  
 
-- **Fixes**
+- **Fixes**  
   - Meissonic with multiple generators  
-  - Kandinsky V2.2 invalid attention processor  
+  - OmniGen with new transformers  
+  - Invalid attention processors  
   - PixArt Sigma Small and Large loading  
-  - TAESD previews with PixArt  
+  - TAESD previews with PixArt and Lumina 2  
+  - VAE Tiling with non-default tile sizes  
+  - Lumina 2 with IPEX  
+  - Nunchaku updated repo  
+  - Double loading of models with custom UNets  
 
 ## Update for 2025-06-02
 

html/reference.json

@@ -239,7 +239,7 @@
   }, 
 
   "VectorSpaceLab OmniGen v1": {
-    "path": "Shitao/OmniGen-v1",
+    "path": "Shitao/OmniGen-v1-diffusers",
     "desc": "OmniGen is a unified image generation model that can generate a wide range of images from multi-modal prompts. It is designed to be simple, flexible and easy to use.",
     "preview": "Shitao--OmniGen-v1.jpg",
     "skip": true

installer.py

@@ -513,7 +513,7 @@ def get_platform():
 def check_python(supported_minors=[], experimental_minors=[], reason=None):
     if supported_minors is None or len(supported_minors) == 0:
         supported_minors = [9, 10, 11, 12]
-        experimental_minors = []
+        experimental_minors = [13]
     t_start = time.time()
     if args.quick:
         return
@@ -546,7 +546,7 @@ def check_diffusers():
     t_start = time.time()
     if args.skip_all or args.skip_git or args.experimental:
         return
-    sha = '6508da6f06a0da1054ae6a808d0025c04b70f0e8' # diffusers commit hash
+    sha = '8adc6003ba4dbf5b61bb4f1ce571e9e55e145a99' # 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 ''
@@ -633,7 +633,7 @@ def install_rocm_zluda():
     log.info(msg)
 
     if sys.platform == "win32": # TODO install: enable ROCm for windows when available
-        check_python(supported_minors=[10, 11], reason='ZLUDA backend requires Python 3.10 or 3.11')
+        #check_python(supported_minors=[9, 10, 11, 12, 13], reason='ZLUDA backend requires a Python version between 3.9 and 3.13')
 
         if args.device_id is not None:
             if os.environ.get('HIP_VISIBLE_DEVICES', None) is not None:
@@ -655,7 +655,7 @@ def install_rocm_zluda():
         if error is None:
             try:
                 zluda_installer.load()
-                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1 torchvision --index-url https://download.pytorch.org/whl/cu118')
+                torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1+cu118 torchvision==0.22.1+cu118 --index-url https://download.pytorch.org/whl/cu118')
             except Exception as e:
                 error = e
                 log.warning(f'Failed to load ZLUDA: {e}')
@@ -663,13 +663,15 @@ def install_rocm_zluda():
             log.info('Using CPU-only torch')
             torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision')
     else:
-        check_python(supported_minors=[9, 10, 11, 12], reason='ROCm backend requires a Python version between 3.9 and 3.12')
+        #check_python(supported_minors=[9, 10, 11, 12, 13], reason='ROCm backend requires a Python version between 3.9 and 3.13')
 
         if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", None) is None:
             os.environ.setdefault('TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL', '1')
 
         if args.use_nightly:
-            if rocm.version is None or float(rocm.version) >= 6.3: # assume the latest if version check fails
+            if rocm.version is None or float(rocm.version) >= 6.4: # assume the latest if version check fails
+                torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.4')
+            elif rocm.version == "6.3":
                 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.3')
             else: # oldest rocm version on nightly is 6.2.4
                 torch_command = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/rocm6.2.4')
@@ -696,22 +698,21 @@ def install_rocm_zluda():
             log.debug(f'ROCm hipBLASLt: arch={device.name} available={device.blaslt_supported}')
             rocm.set_blaslt_enabled(device.blaslt_supported)
 
-    if device is None:
-        log.debug('ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped')
+    if device is None or os.environ.get("HSA_OVERRIDE_GFX_VERSION", None) is not None:
+        log.info(f'ROCm: HSA_OVERRIDE_GFX_VERSION auto config skipped: device={device.name if device is not None else None} version={os.environ.get("HSA_OVERRIDE_GFX_VERSION", None)}')
     else:
         gfx_ver = device.get_gfx_version()
         if gfx_ver is not None:
             os.environ.setdefault('HSA_OVERRIDE_GFX_VERSION', gfx_ver)
-        else:
-            log.warning(f'ROCm: device={device.name} could not auto-detect HSA version')
+            log.info(f'ROCm: HSA_OVERRIDE_GFX_VERSION config overridden: device={device.name} version={os.environ.get("HSA_OVERRIDE_GFX_VERSION", None)}')
 
     ts('amd', t_start)
     return torch_command
 
 
-def install_ipex(torch_command):
+def install_ipex():
     t_start = time.time()
-    check_python(supported_minors=[9, 10, 11, 12], reason='IPEX backend requires a Python version between 3.9 and 3.12')
+    #check_python(supported_minors=[9, 10, 11, 12, 13], reason='IPEX backend requires a Python version between 3.9 and 3.13')
     args.use_ipex = True # pylint: disable=attribute-defined-outside-init
     log.info('IPEX: Intel OneAPI toolkit detected')
 
@@ -742,9 +743,9 @@ def install_ipex(torch_command):
     return torch_command
 
 
-def install_openvino(torch_command):
+def install_openvino():
     t_start = time.time()
-    check_python(supported_minors=[9, 10, 11, 12], reason='OpenVINO backend requires a Python version between 3.9 and 3.12')
+    #check_python(supported_minors=[9, 10, 11, 12, 13], reason='OpenVINO backend requires a Python version between 3.9 and 3.13')
     log.info('OpenVINO: selected')
     if sys.platform == 'darwin':
         torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.7.1 torchvision==0.22.1')
@@ -840,15 +841,15 @@ def check_torch():
         elif is_rocm_available and (args.use_rocm or args.use_zluda): # prioritize rocm
             torch_command = install_rocm_zluda()
         elif allow_ipex and args.use_ipex: # prioritize ipex
-            torch_command = install_ipex(torch_command)
+            torch_command = install_ipex()
         elif allow_openvino and args.use_openvino: # prioritize openvino
-            torch_command = install_openvino(torch_command)
+            torch_command = install_openvino()
         elif is_cuda_available:
             torch_command = install_cuda()
         elif is_rocm_available:
             torch_command = install_rocm_zluda()
         elif is_ipex_available:
-            torch_command = install_ipex(torch_command)
+            torch_command = install_ipex()
         else:
             machine = platform.machine()
             if sys.platform == 'darwin':
@@ -867,6 +868,7 @@ def check_torch():
     if 'torch' in torch_command and not args.version:
         if not installed('torch'):
             log.info(f'Torch: download and install in progress... cmd="{torch_command}"')
+            install('--upgrade pip', 'pip', reinstall=True) # pytorch rocm is too large for older pip
         install(torch_command, 'torch torchvision', quiet=True)
     else:
         try:
@@ -1155,8 +1157,8 @@ def ensure_base_requirements():
 def install_optional():
     t_start = time.time()
     log.info('Installing optional requirements...')
-    install('basicsr')
-    install('gfpgan')
+    install('git+https://github.com/Disty0/BasicSR@2b6a12c28e0c81bfb13b7e984144f0b0f5461484', 'basicsr')
+    install('git+https://github.com/Disty0/GFPGAN@09b1190eabbc77e5f15c61fa7c38a2064b403e20', 'gfpgan')
     install('clean-fid')
     install('pillow-jxl-plugin==1.3.3', ignore=True)
     install('optimum-quanto==0.2.7', ignore=True)
@@ -1188,6 +1190,16 @@ def install_requirements():
         pr.enable()
     if args.skip_requirements and not args.requirements:
         return
+    if int(sys.version_info.minor) >= 13:
+        install('audioop-lts')
+        # gcc 15 patch
+        backup_cmake_policy = os.environ.get('CMAKE_POLICY_VERSION_MINIMUM', None)
+        backup_cxxflags = os.environ.get('CXXFLAGS', None)
+        os.environ.setdefault('CMAKE_POLICY_VERSION_MINIMUM', '3.5')
+        os.environ.setdefault('CXXFLAGS', '-include cstdint')
+        install('git+https://github.com/google/sentencepiece#subdirectory=python', 'sentencepiece')
+        os.environ.setdefault('CMAKE_POLICY_VERSION_MINIMUM', backup_cmake_policy)
+        os.environ.setdefault('CXXFLAGS', backup_cxxflags)
     if not installed('diffusers', quiet=True): # diffusers are not installed, so run initial installation
         global quick_allowed # pylint: disable=global-statement
         quick_allowed = False

javascript/gallery.js

@@ -131,7 +131,7 @@ class GalleryFile extends HTMLElement {
     }
     const ext = this.name.split('.').pop().toLowerCase();
     if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) {
-      console.error(`gallery: type=${ext} file=${this.name} unsupported`);
+      // 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
@@ -286,7 +286,7 @@ async function gallerySearch(evt) {
 
 const findDuplicates = (arr, key) => {
   const map = new Map();
-  return arr.filter(item => {
+  return arr.filter((item) => {
     const value = item[key];
     if (map.has(value)) return true;
     map.set(value, true);

javascript/ui.js

@@ -60,6 +60,17 @@ function selected_gallery_index() {
   return result;
 }
 
+function selected_gallery_files() {
+  let allImages = [];
+  try {
+    let allCurrentButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnail-item.thumbnail-small');
+    if (allCurrentButtons.length === 0) allCurrentButtons = gradioApp().querySelectorAll('.gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small');
+    allImages = Array.from(allCurrentButtons).map((v) => v.querySelector('img')?.src);
+  } catch { /**/ }
+  const selectedIndex = selected_gallery_index();
+  return [allImages, selectedIndex];
+}
+
 function extract_image_from_gallery(gallery) {
   if (gallery.length === 0) return [null];
   if (gallery.length === 1) return [gallery[0]];

modules/api/api.py

@@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request
 from fastapi.security import HTTPBasic, HTTPBasicCredentials
 from fastapi.exceptions import HTTPException
 from modules import errors, shared, postprocessing
-from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, docs
+from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, loras, docs
 
 
 errors.install()
@@ -78,6 +78,7 @@ class Api:
         self.add_api_route("/sdapi/v1/samplers", endpoints.get_samplers, methods=["GET"], response_model=List[models.ItemSampler])
         self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler])
         self.add_api_route("/sdapi/v1/sd-models", endpoints.get_sd_models, methods=["GET"], response_model=List[models.ItemModel])
+        self.add_api_route("/sdapi/v1/controlnets", endpoints.get_controlnets, methods=["GET"], response_model=List[str])
         self.add_api_route("/sdapi/v1/hypernetworks", endpoints.get_hypernetworks, methods=["GET"], response_model=List[models.ItemHypernetwork])
         self.add_api_route("/sdapi/v1/face-restorers", endpoints.get_detailers, methods=["GET"], response_model=List[models.ItemDetailer])
         self.add_api_route("/sdapi/v1/prompt-styles", endpoints.get_prompt_styles, methods=["GET"], response_model=List[models.ItemStyle])
@@ -100,8 +101,9 @@ class Api:
 
         # lora api
         if shared.native:
-            self.add_api_route("/sdapi/v1/loras", endpoints.get_loras, methods=["GET"], response_model=List[dict])
-            self.add_api_route("/sdapi/v1/refresh-loras", endpoints.post_refresh_loras, methods=["POST"])
+            self.add_api_route("/sdapi/v1/lora", loras.get_lora, methods=["GET"], response_model=dict)
+            self.add_api_route("/sdapi/v1/loras", loras.get_loras, methods=["GET"], response_model=List[dict])
+            self.add_api_route("/sdapi/v1/refresh-loras", loras.post_refresh_loras, methods=["POST"])
 
         # gallery api
         gallery.register_api(self.app)

modules/api/endpoints.py

@@ -23,6 +23,10 @@ def get_sd_models():
         checkpoints.append({"title": v.title, "model_name": v.name, "filename": v.filename, "type": v.type, "hash": v.shorthash, "sha256": v.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(v)})
     return checkpoints
 
+def get_controlnets(model_type: Optional[str] = None):
+    from modules.control.units.controlnet import api_list_models
+    return api_list_models(model_type)
+
 def get_hypernetworks():
     return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks]
 
@@ -43,12 +47,6 @@ def get_embeddings():
 
     return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)}
 
-def get_loras():
-    from modules.lora import network, lora_load
-    def create_lora_json(obj: network.NetworkOnDisk):
-        return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
-    return [create_lora_json(obj) for obj in lora_load.available_networks.values()]
-
 def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin
     res = []
     for pg in shared.extra_networks:
@@ -158,10 +156,6 @@ def post_refresh_vae():
     shared.refresh_vaes()
     return {}
 
-def post_refresh_loras():
-    from modules.lora import lora_load
-    return lora_load.list_available_networks()
-
 def get_extensions_list():
     from modules import extensions
     extensions.list_extensions()

modules/api/loras.py

@@ -0,0 +1,21 @@
+from fastapi.exceptions import HTTPException
+
+
+def get_lora(lora: str) -> dict:
+    from modules.lora import lora_load
+    if lora not in lora_load.available_networks:
+        raise HTTPException(status_code=404, detail=f"Lora '{lora}' not found")
+    obj = lora_load.available_networks[lora]
+    obj.info = obj.get_info()
+    obj.desc = obj.get_desc()
+    return obj.__dict__
+
+def get_loras():
+    from modules.lora import network, lora_load
+    def create_lora_json(obj: network.NetworkOnDisk):
+        return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata }
+    return [create_lora_json(obj) for obj in lora_load.available_networks.values()]
+
+def post_refresh_loras():
+    from modules.lora import lora_load
+    return lora_load.list_available_networks()

modules/control/units/controlnet.py

@@ -137,6 +137,23 @@ def find_models():
 
 find_models()
 
+
+def api_list_models(model_type: str = None):
+    import modules.shared
+    model_type = model_type or modules.shared.sd_model_type
+    model_list = []
+    if model_type == 'sd' or model_type == 'all':
+        model_list += list(predefined_sd15)
+    if model_type == 'sdxl' or model_type == 'all':
+        model_list += list(predefined_sdxl)
+    if model_type == 'f1' or model_type == 'all':
+        model_list += list(predefined_f1)
+    if model_type == 'sd3' or model_type == 'all':
+        model_list += list(predefined_sd3)
+    model_list += sorted(find_models())
+    return model_list
+
+
 def list_models(refresh=False):
     import modules.shared
     global models # pylint: disable=global-statement

modules/devices.py

@@ -426,8 +426,6 @@ def override_ipex_math():
 
 def set_sdpa_params():
     try:
-        if opts.cross_attention_optimization != "Scaled-Dot-Product":
-            return
         try:
             global sdpa_original # pylint: disable=global-statement
             if sdpa_original is not None:
@@ -667,6 +665,6 @@ def normalize_device(dev):
 
 
 def same_device(d1, d2):
-    if d1.type != d2.type:
+    if torch.device(d1).type != torch.device(d2).type:
         return False
     return normalize_device(d1) == normalize_device(d2)

modules/generation_parameters_copypaste.py

@@ -3,7 +3,7 @@ import io
 import os
 from PIL import Image
 import gradio as gr
-from modules.paths import data_path
+from modules.paths import params_path
 from modules import shared, gr_tempdir, script_callbacks, images
 from modules.infotext import parse, mapping, quote, unquote # pylint: disable=unused-import
 
@@ -223,9 +223,8 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp
 
     def paste_func(prompt):
         if prompt is None or len(prompt.strip()) == 0:
-            filename = os.path.join(data_path, "params.txt")
-            if os.path.exists(filename):
-                with open(filename, "r", encoding="utf8") as file:
+            if os.path.exists(params_path):
+                with open(params_path, "r", encoding="utf8") as file:
                     prompt = file.read()
                 shared.log.debug(f'Prompt parse: type="params" prompt="{prompt}"')
             else:

modules/gr_tempdir.py

@@ -74,8 +74,9 @@ def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint:
         shared.state.image_history += 1
     params = ', '.join([f'{k}: {v}' for k, v in img.info.items()])
     params = params[12:] if params.startswith('parameters: ') else params
-    with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
-        file.write(params)
+    if len(params) > 2:
+        with open(paths.params_path, "w", encoding="utf8") as file:
+            file.write(params)
     return name
 
 

modules/hidiffusion/__init__.py

@@ -41,5 +41,5 @@ def apply(p, model_type):
 
 def unapply():
     pipe = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
-    if hasattr(pipe, 'unet'):
+    if hasattr(pipe, 'unet') and pipe.unet is not None:
         hidiffusion.remove_hidiffusion(pipe)

modules/images.py

@@ -19,6 +19,7 @@ from modules.video import save_video # pylint: disable=unused-import
 
 
 debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None
+debug_save = errors.log.trace if os.environ.get('SD_SAVE_DEBUG', None) is not None else lambda *args, **kwargs: None
 try:
     from pi_heif import register_heif_opener
     register_heif_opener()
@@ -26,7 +27,6 @@ except Exception:
     pass
 
 
-
 def sanitize_filename_part(text, replace_spaces=True):
     if text is None:
         return None
@@ -47,8 +47,9 @@ def atomically_save_image():
     while True:
         image, filename, extension, params, exifinfo, filename_txt = save_queue.get()
         shared.state.image_history += 1
-        with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file:
-            file.write(exifinfo)
+        if len(exifinfo) > 2:
+            with open(paths.params_path, "w", encoding="utf8") as file:
+                file.write(exifinfo)
         fn = filename + extension
         filename = filename.strip()
         if extension[0] != '.': # add dot if missing
@@ -73,6 +74,7 @@ def atomically_save_image():
             pnginfo_data = PngImagePlugin.PngInfo()
             for k, v in params.pnginfo.items():
                 pnginfo_data.add_text(k, str(v))
+            debug_save(f'Save pnginfo: {params.pnginfo.items()}')
             save_args = { 'compress_level': 6, 'pnginfo': pnginfo_data if shared.opts.image_metadata else None }
         elif image_format == 'JPEG':
             if image.mode == 'RGBA':
@@ -82,12 +84,14 @@ def atomically_save_image():
                 image = image.point(lambda p: p * 0.0038910505836576).convert("L")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         elif image_format == 'WEBP':
             if image.mode == 'I;16':
                 image = image.point(lambda p: p * 0.0038910505836576).convert("RGB")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         elif image_format == 'JXL':
             if image.mode == 'I;16':
@@ -96,10 +100,12 @@ def atomically_save_image():
                 image = image.convert("RGBA")
             save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless }
             if shared.opts.image_metadata:
+                debug_save(f'Save exif: {exifinfo}')
                 save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } })
         else:
             save_args = { 'quality': shared.opts.jpeg_quality }
         try:
+            debug_save(f'Save args: {save_args}')
             image.save(fn, format=image_format, **save_args)
         except Exception as e:
             shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}')

modules/intel/ipex/__init__.py

@@ -39,7 +39,6 @@ def ipex_init(): # pylint: disable=too-many-statements
             torch.cuda.is_available = torch.xpu.is_available
             torch.cuda.is_initialized = torch.xpu.is_initialized
             torch.cuda.is_current_stream_capturing = lambda: False
-            torch.cuda.set_device = torch.xpu.set_device
             torch.cuda.stream = torch.xpu.stream
             torch.cuda.Event = torch.xpu.Event
             torch.cuda.Stream = torch.xpu.Stream

modules/intel/ipex/diffusers.py

@@ -81,14 +81,46 @@ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos, output_type="np"):
     return emb
 
 
+def apply_rotary_emb(x, freqs_cis, use_real: bool = True, use_real_unbind_dim: int = -1):
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+        return out
+    else:
+        # used for lumina
+        # force cpu with Alchemist
+        x_rotated = torch.view_as_complex(x.to("cpu").float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.to("cpu").unsqueeze(2)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+        return x_out.type_as(x).to(x.device)
+
+
 def ipex_diffusers(device_supports_fp64=False):
     diffusers.utils.torch_utils.fourier_filter = fourier_filter
     if not device_supports_fp64:
         # get around lazy imports
+        from diffusers.models import embeddings as diffusers_embeddings # pylint: disable=import-error, unused-import # noqa: F401
         from diffusers.models import transformers as diffusers_transformers # pylint: disable=import-error, unused-import # noqa: F401
         from diffusers.models import controlnets as diffusers_controlnets # pylint: disable=import-error, unused-import # noqa: F401
         diffusers.models.embeddings.get_1d_sincos_pos_embed_from_grid = get_1d_sincos_pos_embed_from_grid
         diffusers.models.embeddings.FluxPosEmbed = FluxPosEmbed
+        diffusers.models.embeddings.apply_rotary_emb = apply_rotary_emb
         diffusers.models.transformers.transformer_flux.FluxPosEmbed = FluxPosEmbed
+        diffusers.models.transformers.transformer_lumina2.apply_rotary_emb = apply_rotary_emb
         diffusers.models.controlnets.controlnet_flux.FluxPosEmbed = FluxPosEmbed
         diffusers.models.transformers.transformer_hidream_image.rope = hidream_rope

modules/intel/ipex/hijacks.py

@@ -254,8 +254,15 @@ torch.Tensor.original_Tensor_to = torch.Tensor.to
 @wraps(torch.Tensor.to)
 def Tensor_to(self, device=None, *args, **kwargs):
     if check_cuda(device):
+        if not device_supports_fp64 and kwargs.get("dtype", None) == torch.float64:
+            kwargs["dtype"] = torch.float32
         return self.original_Tensor_to(return_xpu(device), *args, **kwargs)
     else:
+        if not device_supports_fp64:
+            if kwargs.get("dtype", None) == torch.float64 and ((device is None and self.device.type == "xpu") or (device is not None and torch.device(device).type == "xpu")):
+                kwargs["dtype"] = torch.float32
+            elif device == torch.float64 and self.device.type == "xpu":
+                device = torch.float32
         return self.original_Tensor_to(device, *args, **kwargs)
 
 original_Tensor_cuda = torch.Tensor.cuda
@@ -379,6 +386,12 @@ def torch_cuda_device(device):
     else:
         return torch.xpu.device(device)
 
+@wraps(torch.cuda.set_device)
+def torch_cuda_set_device(device):
+    if check_cuda(device):
+        torch.xpu.set_device(return_xpu(device))
+    else:
+        torch.xpu.set_device(device)
 
 # torch.Generator has to be a class for isinstance checks
 original_torch_Generator = torch.Generator
@@ -412,6 +425,7 @@ def ipex_hijacks():
     torch.load = torch_load
     torch.cuda.synchronize = torch_cuda_synchronize
     torch.cuda.device = torch_cuda_device
+    torch.cuda.set_device = torch_cuda_set_device
 
     torch.Generator = torch_Generator
     torch._C.Generator = torch_Generator

modules/interrogate/joycaption.py

@@ -58,9 +58,12 @@ opts = JoyOptions()
 
 
 @torch.no_grad()
-def predict(question: str, image):
+def predict(question: str, image, vqa_model: str = None) -> str:
     global llava_model, processor # pylint: disable=global-statement
     opts.max_new_tokens = shared.opts.interrogate_vlm_max_length
+    if vqa_model is not None and opts.repo != vqa_model:
+        opts.repo = vqa_model
+        llava_model = None
     if llava_model is None:
         shared.log.info(f'Interrogate: type=vlm model="JoyCaption" {str(opts)}')
         processor = AutoProcessor.from_pretrained(opts.repo)

modules/interrogate/vqa.py

@@ -38,7 +38,8 @@ vlm_models = {
     "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B",
     "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B",
     "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB
-    "JoyCaption": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB
+    "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", # 17.4GB
     "JoyTag": "fancyfeast/joytag", # 0.7GB
     "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B",
     "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B",
@@ -583,7 +584,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image:
             answer = joytag.predict(image)
         elif 'joycaption' in vqa_model.lower():
             from modules.interrogate import joycaption
-            answer = joycaption.predict(question, image)
+            answer = joycaption.predict(question, image, vqa_model)
         elif 'deepseek' in vqa_model.lower():
             from modules.interrogate import deepseek
             answer = deepseek.predict(question, image, vqa_model)

modules/ipadapter.py

@@ -146,7 +146,7 @@ def unapply(pipe, unload: bool = False): # pylint: disable=arguments-differ
             if unload:
                 shared.log.debug('IP adapter unload')
                 pipe.unload_ip_adapter()
-        if hasattr(pipe, 'unet'):
+        if hasattr(pipe, 'unet') and pipe.unet is not None:
             module = pipe.unet
         elif hasattr(pipe, 'transformer'):
             module = pipe.transformer

modules/lora/lora_apply.py

@@ -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 hasattr(self, "sdnq_decompressor"):
-                        self.sdnq_decompressor_backup = self.sdnq_decompressor.to(devices.cpu)
+                    if hasattr(self, "sdnq_dequantizer"):
+                        self.sdnq_dequantizer_backup = self.sdnq_dequantizer.to(devices.cpu)
 
         if bias_backup is None:
             if getattr(self, 'bias', None) is not None:
@@ -79,8 +79,10 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn.
             continue
         try:
             t0 = time.time()
-            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)
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul)
+            elif hasattr(self, "sdnq_dequantizer"):
+                weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.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)
@@ -136,25 +138,28 @@ 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 hasattr(self, "sdnq_decompressor"):
+    elif not bias and hasattr(self, "sdnq_dequantizer"):
         try:
-            from modules.model_quant_sdnq import sdnq_quantize_layer
-            if hasattr(self, "sdnq_decompressor_backup"):
-                sdnq_decompressor = self.sdnq_decompressor_backup.to(devices.device)
+            from modules.sdnq import sdnq_quantize_layer
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                sdnq_dequantizer = self.sdnq_dequantizer_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)
+                sdnq_dequantizer = self.sdnq_dequantizer.to(devices.device)
+            dequant_weight = sdnq_dequantizer(model_weights.to(devices.device), skip_quantized_matmul=sdnq_dequantizer.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.sdnq_decompressor = None
+            self.sdnq_dequantizer = None
             self = sdnq_quantize_layer(
                 self,
-                sdnq_decompressor.weights_dtype,
+                sdnq_dequantizer.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,
                 use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+                dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+                quantization_device=devices.device,
+                return_device=device,
                 param_name=getattr(self, 'network_layer_name', None),
             )
             self = self.to(device)
@@ -223,8 +228,9 @@ 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, "sdnq_decompressor_backup"):
-                self.sdnq_decompressor = self.sdnq_decompressor_backup.to(device)
+            if hasattr(self, "sdnq_dequantizer_backup"):
+                self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device)
+                del self.sdnq_dequantizer_backup
 
     if bias_backup is not None:
         self.bias = None

modules/lora/network.py

@@ -90,6 +90,27 @@ class NetworkOnDisk:
         if not self.hash:
             self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '')
 
+    def get_info(self):
+        data = {}
+        if shared.cmd_opts.no_metadata:
+            return data
+        if self.filename is not None:
+            fn = os.path.splitext(self.filename)[0] + '.json'
+            if os.path.exists(fn):
+                data = shared.readfile(fn, silent=True)
+                if type(data) is list:
+                    data = data[0]
+        return data
+
+    def get_desc(self):
+        if shared.cmd_opts.no_metadata:
+            return None
+        if self.filename is not None:
+            fn = os.path.splitext(self.filename)[0] + '.txt'
+            if os.path.exists(fn):
+                return shared.readfile(fn, silent=True)
+        return None
+
     def get_alias(self):
         if shared.opts.lora_preferred_name == "filename":
             return self.name
@@ -124,7 +145,10 @@ class NetworkModule:
         self.sd_key = weights.sd_key
         self.sd_module = weights.sd_module
         if hasattr(self.sd_module, 'weight'):
-            self.shape = self.sd_module.weight.shape
+            if hasattr(self.sd_module, "sdnq_dequantizer"):
+                self.shape = self.sd_module.sdnq_dequantizer.original_shape
+            else:
+                self.shape = self.sd_module.weight.shape
         self.dim = None
         self.bias = weights.w.get("bias")
         self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None

modules/model_flux.py

@@ -14,7 +14,6 @@ debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None el
 def load_flux_quanto(checkpoint_info):
     transformer, text_encoder_2 = None, None
     quanto = model_quant.load_quanto('Load model: type=FLUX')
-    quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
 
     if isinstance(checkpoint_info, str):
         repo_path = checkpoint_info
@@ -36,11 +35,12 @@ def load_flux_quanto(checkpoint_info):
         quanto.requantize(transformer, state_dict, quantization_map, device=torch.device("cpu"))
         if shared.opts.diffusers_eval:
             transformer.eval()
-        if transformer.dtype != devices.dtype:
+        transformer_dtype = transformer.dtype
+        if transformer_dtype != devices.dtype:
             try:
                 transformer = transformer.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer.dtype}")
+                shared.log.error(f"Load model: type=FLUX Failed to cast transformer to {devices.dtype}, set dtype to {transformer_dtype}")
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto transformer: {e}")
         if debug:
@@ -63,11 +63,12 @@ def load_flux_quanto(checkpoint_info):
         quanto.requantize(text_encoder_2, state_dict, quantization_map, device=torch.device("cpu"))
         if shared.opts.diffusers_eval:
             text_encoder_2.eval()
-        if text_encoder_2.dtype != devices.dtype:
+        text_encoder_2_dtype = text_encoder_2.dtype
+        if text_encoder_2_dtype != devices.dtype:
             try:
                 text_encoder_2 = text_encoder_2.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2.dtype}")
+                shared.log.error(f"Load model: type=FLUX Failed to cast text encoder to {devices.dtype}, set dtype to {text_encoder_2_dtype}")
     except Exception as e:
         shared.log.error(f"Load model: type=FLUX failed to load Quanto text encoder: {e}")
         if debug:
@@ -114,11 +115,11 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant):
             nunchaku_precision = nunchaku.utils.get_precision()
             nunchaku_repo = None
             if 'dev' in repo_id:
-                nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-flux.1-dev"
+                nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-dev/svdq-{nunchaku_precision}_r32-flux.1-dev.safetensors"
             elif 'schnell' in repo_id:
-                nunchaku_repo = f"mit-han-lab/svdq-{nunchaku_precision}-flux.1-schnell"
+                nunchaku_repo = f"mit-han-lab/nunchaku-flux.1-schnell/svdq-{nunchaku_precision}_r32-flux.1-schnell.safetensors"
             elif 'shuttle' in repo_id:
-                nunchaku_repo = 'mit-han-lab/svdq-fp4-shuttle-jaguar'
+                nunchaku_repo = f"mit-han-lab/nunchaku-shuttle-jaguar/svdq-{nunchaku_precision}_r32-shuttle-jaguar.safetensors"
             else:
                 shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported')
             if nunchaku_repo is not None:
@@ -134,7 +135,7 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant):
         if 'text_encoder_2' not in kwargs and model_quant.check_nunchaku('TE'):
             import nunchaku
             nunchaku_precision = nunchaku.utils.get_precision()
-            nunchaku_repo = 'mit-han-lab/svdq-flux.1-t5'
+            nunchaku_repo = 'mit-han-lab/nunchaku-t5/awq-int4-flux.1-t5xxl.safetensors'
             shared.log.debug(f'Load module: quant=Nunchaku module=t5 repo="{nunchaku_repo}" precision={nunchaku_precision}')
             kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
         elif 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'):

modules/model_lumina.py

@@ -1,5 +1,10 @@
+import os
 import transformers
 import diffusers
+from huggingface_hub import repo_exists
+from modules import errors, shared, sd_unet, sd_hijack_te
+
+debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
 
 
 def load_lumina(_checkpoint_info, diffusers_load_config={}):
@@ -17,29 +22,72 @@ def load_lumina(_checkpoint_info, diffusers_load_config={}):
 
 def load_lumina2(checkpoint_info, diffusers_load_config={}):
     from modules import shared, devices, sd_models, model_quant
+    transformer, text_encoder, vae = None, None, None
     repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    if os.path.isdir(checkpoint_info.filename) and not repo_exists(repo_id):
+        repo_id = checkpoint_info.filename
+
+    if shared.opts.teacache_enabled:
+        from modules import teacache
+        shared.log.debug(f'Transformers cache: type=teacache patch=forward cls={diffusers.Lumina2Transformer2DModel.__name__}')
+        diffusers.Lumina2Transformer2DModel.forward = teacache.teacache_lumina2_forward # patch must be done before transformer is loaded
 
     load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
-    transformer = diffusers.Lumina2Transformer2DModel.from_pretrained(
-        repo_id,
-        subfolder="transformer",
-        cache_dir=shared.opts.hfcache_dir,
-        **load_config,
-        **quant_config,
-    )
+    if shared.opts.sd_unet != 'Default':
+        try:
+            debug(f'Load model: type=Lumina2 unet="{shared.opts.sd_unet}"')
+            transformer = diffusers.Lumina2Transformer2DModel.from_single_file(
+                sd_unet.unet_dict[shared.opts.sd_unet],
+                cache_dir=shared.opts.diffusers_dir,
+                **load_config,
+                **quant_config
+            )
+            if transformer is None:
+                shared.opts.sd_unet = 'Default'
+                sd_unet.failed_unet.append(shared.opts.sd_unet)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Lumina2 failed to load UNet: {e}")
+            shared.opts.sd_unet = 'Default'
+            if debug:
+                errors.display(e, 'Lumina2 UNet:')
+
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
+        try:
+            debug(f'Load model: type=Lumina2 vae="{shared.opts.sd_vae}"')
+            from modules import sd_vae
+            # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override')
+            vae_file = sd_vae.vae_dict[shared.opts.sd_vae]
+            if os.path.exists(vae_file):
+                vae_config = os.path.join('configs', 'flux', 'vae', 'config.json')
+                vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
+        except Exception as e:
+            shared.log.error(f"Load model: type=Lumina2 failed to load VAE: {e}")
+            shared.opts.sd_vae = 'Default'
+            if debug:
+                errors.display(e, 'Lumina2 VAE:')
+
+    if transformer is None:
+        transformer = diffusers.Lumina2Transformer2DModel.from_pretrained(
+            repo_id,
+            subfolder="transformer",
+            cache_dir=shared.opts.diffusers_dir,
+            **load_config,
+            **quant_config,
+        )
 
     load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)
     text_encoder = transformers.AutoModel.from_pretrained(
         repo_id,
         subfolder="text_encoder",
-        cache_dir=shared.opts.hfcache_dir,
-        torch_dtype=devices.dtype,
+        cache_dir=shared.opts.diffusers_dir,
         **load_config,
         **quant_config,
     )
 
     load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
-    pipe = diffusers.Lumina2Text2ImgPipeline.from_pretrained(
+    if vae is not None:
+        load_config['vae'] = vae
+    pipe = diffusers.Lumina2Pipeline.from_pretrained(
         repo_id,
         cache_dir=shared.opts.diffusers_dir,
         text_encoder=text_encoder,
@@ -47,5 +95,6 @@ def load_lumina2(checkpoint_info, diffusers_load_config={}):
         **load_config,
     )
 
+    sd_hijack_te.init_hijack(pipe)
     devices.torch_gc(force=True)
     return pipe

modules/model_omnigen.py

@@ -1,25 +1,47 @@
-def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
-    from modules import shared, devices, sd_models, shared_items
-    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+import os
+import diffusers
+from modules import errors, shared, devices, sd_models, model_quant
 
-    # load
-    from modules.omnigen import OmniGenPipeline
-    shared_items.pipelines['OmniGen'] = OmniGenPipeline
-    pipe = OmniGenPipeline.from_pretrained(
-        model_name=repo_id,
-        vae_path='madebyollin/sdxl-vae-fp16-fix',
+debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None
+
+
+def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument
+    repo_id = sd_models.path_to_repo(checkpoint_info.name)
+    vae = None
+
+    if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic':
+        try:
+            debug(f'Load model: type=OmniGen vae="{shared.opts.sd_vae}"')
+            from modules import sd_vae
+            # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override')
+            vae_file = sd_vae.vae_dict[shared.opts.sd_vae]
+            if os.path.exists(vae_file):
+                vae_config = os.path.join('configs', 'sdxl', 'vae', 'config.json')
+                vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config)
+        except Exception as e:
+            shared.log.error(f"Load model: type=OmniGen failed to load VAE: {e}")
+            shared.opts.sd_vae = 'Default'
+            if debug:
+                errors.display(e, 'OmniGen VAE:')
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer')
+    transformer = diffusers.OmniGenTransformer2DModel.from_pretrained(
+        repo_id,
+        subfolder="transformer",
         cache_dir=shared.opts.diffusers_dir,
+        **load_config,
+        **quant_config,
+    )
+
+    load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+    if vae is not None:
+        load_config['vae'] = vae
+    pipe = diffusers.OmniGenPipeline.from_pretrained(
+        repo_id,
+        transformer=transformer,
+        cache_dir=shared.opts.diffusers_dir,
+        **load_config,
     )
 
-    # init
-    pipe.device = devices.device
-    pipe.dtype = devices.dtype
-    pipe.model.device = devices.device
-    pipe.separate_cfg_infer = True
-    pipe.use_kv_cache = False
-    pipe.model.to(device=devices.device, dtype=devices.dtype)
-    if shared.opts.diffusers_eval:
-        pipe.model.eval()
-    pipe.vae.to(devices.device, dtype=devices.dtype)
     devices.torch_gc(force=True)
     return pipe

modules/model_quant.py

@@ -104,23 +104,52 @@ def create_quanto_config(kwargs = None, allow_quanto: bool = True, module: str =
 
 
 def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Model', weights_dtype: str = None):
-    from modules import shared
+    from modules import devices, shared
     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
+            from modules.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
 
+            if weights_dtype is None:
+                if module in {"TE", "LLM"}:
+                    if shared.opts.sdnq_quantize_weights_mode_te == "none":
+                        return kwargs
+                    elif shared.opts.sdnq_quantize_weights_mode_te in {"same as model", "default"}:
+                        weights_dtype = shared.opts.sdnq_quantize_weights_mode
+                    else:
+                        weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+                elif shared.opts.sdnq_quantize_weights_mode == "none":
+                    return kwargs
+                else:
+                    weights_dtype = shared.opts.sdnq_quantize_weights_mode
+
+            if shared.opts.device_map == "gpu":
+                quantization_device = devices.device
+                return_device = devices.device
+            elif shared.opts.diffusers_offload_mode in {"none", "model"}:
+                quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
+                return_device = devices.device
+            elif shared.opts.sdnq_quantize_with_gpu:
+                quantization_device = devices.device
+                return_device = devices.cpu
+            else:
+                quantization_device = None
+                return_device = None
+
             sdnq_config = SDNQConfig(
-                weights_dtype=weights_dtype if weights_dtype is not None else shared.opts.sdnq_quantize_weights_mode,
+                weights_dtype=weights_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,
                 use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+                dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
             )
-            log.debug(f'Quantization: module="{module}" type=sdnq dtype={shared.opts.sdnq_quantize_weights_mode}')
+            log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}')
             if kwargs is None:
                 return sdnq_config
             else:
@@ -303,32 +332,50 @@ def apply_layerwise(sd_model, quiet:bool=False):
 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 modules.model_quant_sdnq import apply_sdnq_to_module
+    from modules.sdnq import apply_sdnq_to_module
 
     model.eval()
-
-    if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}:
-        import torch
-        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 = SDNQ_T5DenseGatedActDense(
-                model.encoder.block[i].layer[1].DenseReluDense,
-                dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16
-            )
-
     backup_embeddings = None
     if hasattr(model, "get_input_embeddings"):
         backup_embeddings = copy.deepcopy(model.get_input_embeddings())
 
+    if shared.opts.sdnq_quantize_weights_mode_te != "default" and op is not None and "text_encoder" in op:
+        weights_dtype = shared.opts.sdnq_quantize_weights_mode_te
+    else:
+        weights_dtype = shared.opts.sdnq_quantize_weights_mode
+
+    if weights_dtype is None or weights_dtype == 'none':
+        return model
+    if debug:
+        log.trace(f'Quantization: type=SDNQ op={op} cls={model.__class__} dtype={weights_dtype} mode{shared.opts.diffusers_offload_mode}')
+
+    if shared.opts.diffusers_offload_mode in {"none", "model"}:
+        quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu
+        return_device = devices.device
+    elif shared.opts.sdnq_quantize_with_gpu:
+        quantization_device = devices.device
+        return_device = getattr(model, "device", devices.cpu)
+    else:
+        quantization_device = None
+        return_device = None
+
+    modules_to_not_convert = getattr(model, "_keep_in_fp32_modules", [])
+    if modules_to_not_convert is None:
+        modules_to_not_convert = []
+
     model = apply_sdnq_to_module(
         model,
-        weights_dtype=shared.opts.sdnq_quantize_weights_mode,
+        weights_dtype=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,
         use_quantized_matmul_conv=shared.opts.sdnq_use_quantized_matmul_conv,
+        dequantize_fp32=shared.opts.sdnq_dequantize_fp32,
+        quantization_device=quantization_device,
+        return_device=return_device,
         param_name=op,
+        modules_to_not_convert=modules_to_not_convert,
     )
     model.quantization_method = 'SDNQ'
 
@@ -362,7 +409,7 @@ def sdnq_quantize_weights(sd_model):
     try:
         t0 = time.time()
         from modules import shared, devices, sd_models
-        log.info(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights}")
+        log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul}  group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32}")
         global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
 
         sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq")

modules/model_te.py

@@ -4,7 +4,6 @@ import torch
 import transformers
 from safetensors.torch import load_file
 from modules import shared, devices, files_cache, errors, model_quant
-from installer import install
 
 
 te_dict = {}
@@ -72,27 +71,32 @@ def load_t5(name=None, cache_dir=None):
     elif 'int8' 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='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not 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 '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)
+        if quantization_config is not 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 'qint4' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int4')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
 
     elif 'qint8' in name.lower():
         model_quant.load_quanto('Load model: type=T5')
         quantization_config = transformers.QuantoConfig(weights='int8')
-        t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype)
+        if quantization_config is not 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 '/' in name:
         shared.log.debug(f'Load model: type=T5 repo={name}')
         quant_config = model_quant.create_config(module='TE')
-        t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
+        if quantization_config is not None:
+            t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config)
 
     else:
         t5 = None

modules/omnigen/__init__.py

@@ -1,4 +0,0 @@
-from .model import OmniGen
-from .processor import OmniGenProcessor
-from .scheduler import OmniGenScheduler
-from .pipeline import OmniGenPipeline

modules/omnigen/model.py

@@ -1,390 +0,0 @@
-# The code is revised from DiT
-import os
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-from safetensors.torch import load_file
-from diffusers.loaders import PeftAdapterMixin
-from huggingface_hub import snapshot_download
-from .transformer import Phi3Config, Phi3Transformer
-
-
-def modulate(x, shift, scale):
-    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
-
-
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        """
-        Create sinusoidal timestep embeddings.
-        :param t: a 1-D Tensor of N indices, one per batch element.
-                          These may be fractional.
-        :param dim: the dimension of the output.
-        :param max_period: controls the minimum frequency of the embeddings.
-        :return: an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        half = dim // 2
-        freqs = torch.exp(
-            -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=t.device)
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t, dtype=torch.float32):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class FinalLayer(nn.Module):
-    """
-    The final layer of DiT.
-    """
-    def __init__(self, hidden_size, patch_size, out_channels):
-        super().__init__()
-        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
-        self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True)
-        self.adaLN_modulation = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(hidden_size, 2 * hidden_size, bias=True)
-        )
-
-    def forward(self, x, c):
-        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
-        x = modulate(self.norm_final(x), shift, scale)
-        x = self.linear(x)
-        return x
-
-
-def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1):
-    """
-    grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or
-    [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
-    """
-    if isinstance(grid_size, int):
-        grid_size = (grid_size, grid_size)
-
-    grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale
-    grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale
-    grid = np.meshgrid(grid_w, grid_h)  # here w goes first
-    grid = np.stack(grid, axis=0)
-
-    grid = grid.reshape([2, 1, grid_size[1], grid_size[0]])
-    pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
-    if cls_token and extra_tokens > 0:
-        pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0)
-    return pos_embed
-
-
-def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
-    assert embed_dim % 2 == 0
-
-    # use half of dimensions to encode grid_h
-    emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0])  # (H*W, D/2)
-    emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1])  # (H*W, D/2)
-
-    emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
-    return emb
-
-
-def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
-    """
-    embed_dim: output dimension for each position
-    pos: a list of positions to be encoded: size (M,)
-    out: (M, D)
-    """
-    assert embed_dim % 2 == 0
-    omega = np.arange(embed_dim // 2, dtype=np.float64)
-    omega /= embed_dim / 2.
-    omega = 1. / 10000**omega  # (D/2,)
-
-    pos = pos.reshape(-1)  # (M,)
-    out = np.einsum('m,d->md', pos, omega)  # (M, D/2), outer product
-
-    emb_sin = np.sin(out) # (M, D/2)
-    emb_cos = np.cos(out) # (M, D/2)
-
-    emb = np.concatenate([emb_sin, emb_cos], axis=1)  # (M, D)
-    return emb
-
-
-class PatchEmbedMR(nn.Module):
-    """ 2D Image to Patch Embedding
-    """
-    def __init__(
-            self,
-            patch_size: int = 2,
-            in_chans: int = 4,
-            embed_dim: int = 768,
-            bias: bool = True,
-    ):
-        super().__init__()
-        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
-
-    def forward(self, x):
-        x = self.proj(x)
-        x = x.flatten(2).transpose(1, 2)  # NCHW -> NLC
-        return x
-
-
-class OmniGen(nn.Module, PeftAdapterMixin):
-    """
-    Diffusion model with a Transformer backbone.
-    """
-    def __init__(
-        self,
-        transformer_config: Phi3Config,
-        patch_size=2,
-        in_channels=4,
-        pe_interpolation: float = 1.0,
-        pos_embed_max_size: int = 192,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-        self.patch_size = patch_size
-        self.pos_embed_max_size = pos_embed_max_size
-        hidden_size = transformer_config.hidden_size
-        self.x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
-        self.input_x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True)
-        self.time_token = TimestepEmbedder(hidden_size)
-        self.t_embedder = TimestepEmbedder(hidden_size)
-        self.pe_interpolation = pe_interpolation
-        pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, interpolation_scale=self.pe_interpolation, base_size=64)
-        self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=True)
-        self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels)
-        self.initialize_weights()
-        self.llm = Phi3Transformer(config=transformer_config)
-        self.llm.config.use_cache = False
-
-    @classmethod
-    def from_pretrained(cls, model_name: str, cache_dir: str=None):
-        if not os.path.exists(os.path.join(model_name, 'model.pt')) and not os.path.exists(os.path.join(model_name, 'model.safetensors')):
-            cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_dir,
-                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
-        config = Phi3Config.from_pretrained(model_name)
-        model = cls(config)
-        if os.path.exists(os.path.join(model_name, 'model.pt')):
-            state_dict = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu')
-        elif os.path.exists(os.path.join(model_name, 'model.safetensors')):
-            state_dict = load_file(os.path.join(model_name, 'model.safetensors'))
-        else:
-            raise ValueError(f"OmniGen: Could not find model file in {model_name}")
-        model.load_state_dict(state_dict)
-        return model
-
-    def initialize_weights(self):
-        assert not hasattr(self, "llama")
-
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-        # Initialize patch_embed like nn.Linear (instead of nn.Conv2d):
-        w = self.x_embedder.proj.weight.data
-        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj.bias, 0)
-
-        w = self.input_x_embedder.proj.weight.data
-        nn.init.xavier_uniform_(w.view([w.shape[0], -1]))
-        nn.init.constant_(self.x_embedder.proj.bias, 0)
-
-
-        # Initialize timestep embedding MLP:
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-        nn.init.normal_(self.time_token.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.time_token.mlp[2].weight, std=0.02)
-
-        # Zero-out output layers:
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
-        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
-        nn.init.constant_(self.final_layer.linear.weight, 0)
-        nn.init.constant_(self.final_layer.linear.bias, 0)
-
-    def unpatchify(self, x, h, w):
-        """
-        x: (N, T, patch_size**2 * C)
-        imgs: (N, H, W, C)
-        """
-        c = self.out_channels
-
-        x = x.reshape(shape=(x.shape[0], h//self.patch_size, w//self.patch_size, self.patch_size, self.patch_size, c))
-        x = torch.einsum('nhwpqc->nchpwq', x)
-        imgs = x.reshape(shape=(x.shape[0], c, h, w))
-        return imgs
-
-
-    def cropped_pos_embed(self, height, width):
-        """Crops positional embeddings for SD3 compatibility."""
-        if self.pos_embed_max_size is None:
-            raise ValueError("`pos_embed_max_size` must be set for cropping.")
-
-        height = height // self.patch_size
-        width = width // self.patch_size
-        if height > self.pos_embed_max_size:
-            raise ValueError(
-                f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
-            )
-        if width > self.pos_embed_max_size:
-            raise ValueError(
-                f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}."
-            )
-
-        top = (self.pos_embed_max_size - height) // 2
-        left = (self.pos_embed_max_size - width) // 2
-        spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1)
-        spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :]
-        spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1])
-        return spatial_pos_embed
-
-
-    def patch_multiple_resolutions(self, latents, padding_latent=None, is_input_images:bool=False):
-        if isinstance(latents, list):
-            return_list = False
-            if padding_latent is None:
-                padding_latent = [None] * len(latents)
-                return_list = True
-            patched_latents, num_tokens, shapes = [], [], []
-            for latent, padding in zip(latents, padding_latent):
-                height, width = latent.shape[-2:]
-                if is_input_images:
-                    latent = self.input_x_embedder(latent)
-                else:
-                    latent = self.x_embedder(latent)
-                pos_embed = self.cropped_pos_embed(height, width)
-                latent = latent + pos_embed
-                if padding is not None:
-                    latent = torch.cat([latent, padding], dim=-2)
-                patched_latents.append(latent)
-
-                num_tokens.append(pos_embed.size(1))
-                shapes.append([height, width])
-            if not return_list:
-                latents = torch.cat(patched_latents, dim=0)
-            else:
-                latents = patched_latents
-        else:
-            height, width = latents.shape[-2:]
-            if is_input_images:
-                latents = self.input_x_embedder(latents)
-            else:
-                latents = self.x_embedder(latents)
-            pos_embed = self.cropped_pos_embed(height, width)
-            latents = latents + pos_embed
-            num_tokens = latents.size(1)
-            shapes = [height, width]
-        return latents, num_tokens, shapes
-
-    def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True):
-        input_is_list = isinstance(x, list)
-        x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent)
-        time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1)
-        if input_img_latents is not None:
-            input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True)
-        if input_ids is not None:
-            condition_embeds = self.llm.embed_tokens(input_ids).clone()
-            input_img_inx = 0
-            for b_inx in input_image_sizes.keys():
-                for start_inx, end_inx in input_image_sizes[b_inx]:
-                    condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx]
-                    input_img_inx += 1
-            if input_img_latents is not None:
-                assert input_img_inx == len(input_latents)
-
-            input_emb = torch.cat([condition_embeds, time_token, x], dim=1)
-        else:
-            input_emb = torch.cat([time_token, x], dim=1)
-        output = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values)
-        output, past_key_values = output.last_hidden_state, output.past_key_values
-        if input_is_list:
-            image_embedding = output[:, -max(num_tokens):]
-            time_emb = self.t_embedder(timestep, dtype=x.dtype)
-            x = self.final_layer(image_embedding, time_emb)
-            latents = []
-            for i in range(x.size(0)):
-                latent = x[i:i+1, :num_tokens[i]]
-                latent = self.unpatchify(latent, shapes[i][0], shapes[i][1])
-                latents.append(latent)
-        else:
-            image_embedding = output[:, -num_tokens:]
-            time_emb = self.t_embedder(timestep, dtype=x.dtype)
-            x = self.final_layer(image_embedding, time_emb)
-            latents = self.unpatchify(x, shapes[0], shapes[1])
-
-        if return_past_key_values:
-            return latents, past_key_values
-        return latents
-
-    @torch.no_grad()
-    def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache):
-        """
-        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """
-        self.llm.config.use_cache = use_kv_cache
-        model_out, past_key_values = self.forward(x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, past_key_values=past_key_values, return_past_key_values=True)
-        if use_img_cfg:
-            cond, uncond, img_cond = torch.split(model_out, len(model_out) // 3, dim=0)
-            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
-            model_out = [cond, cond, cond]
-        else:
-            cond, uncond = torch.split(model_out, len(model_out) // 2, dim=0)
-            cond = uncond + cfg_scale * (cond - uncond)
-            model_out = [cond, cond]
-        return torch.cat(model_out, dim=0), past_key_values
-
-
-    @torch.no_grad()
-    def forward_with_separate_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, return_past_key_values=True):
-        """
-        Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance.
-        """
-        self.llm.config.use_cache = use_kv_cache
-        if past_key_values is None:
-            past_key_values = [None] * len(attention_mask)
-
-        x = torch.split(x, len(x) // len(attention_mask), dim=0)
-        timestep = timestep.to(x[0].dtype)
-        timestep = torch.split(timestep, len(timestep) // len(input_ids), dim=0)
-
-        model_out, pask_key_values = [], []
-        for i in range(len(input_ids)):
-            temp_out, temp_pask_key_values = self.forward(x[i], timestep[i], input_ids[i], input_img_latents[i], input_image_sizes[i], attention_mask[i], position_ids[i], past_key_values[i])
-            model_out.append(temp_out)
-            pask_key_values.append(temp_pask_key_values)
-
-        if len(model_out) == 3:
-            cond, uncond, img_cond = model_out
-            cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond)
-            model_out = [cond, cond, cond]
-        elif len(model_out) == 2:
-            cond, uncond = model_out
-            cond = uncond + cfg_scale * (cond - uncond)
-            model_out = [cond, cond]
-        else:
-            return model_out[0]
-        return torch.cat(model_out, dim=0), pask_key_values

modules/omnigen/pipeline.py

@@ -1,219 +0,0 @@
-import os
-from typing import List, Union
-from PIL import Image
-import torch
-from huggingface_hub import snapshot_download
-from peft import PeftModel
-from diffusers.models import AutoencoderKL
-from diffusers.utils import replace_example_docstring
-from .model import OmniGen
-from .processor import OmniGenProcessor
-from .scheduler import OmniGenScheduler
-
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> from OmniGen import OmniGenPipeline
-        >>> pipe = FluxControlNetPipeline.from_pretrained(
-        ...     base_model
-        ... )
-        >>> prompt = "A woman holds a bouquet of flowers and faces the camera"
-        >>> image = pipe(
-        ...     prompt,
-        ...     guidance_scale=3.0,
-        ...     num_inference_steps=50,
-        ... ).images[0]
-        >>> image.save("t2i.png")
-        ```
-"""
-
-
-class OmniGenPipeline():
-    def __init__(
-        self,
-        vae: AutoencoderKL,
-        model: OmniGen,
-        processor: OmniGenProcessor,
-    ):
-        super().__init__()
-        self.vae = vae
-        self.model = model
-        self.processor = processor
-        self.device = None
-        self.dtype: None
-        self.separate_cfg_infer: bool = True
-        self.use_kv_cache: bool = False
-        # omnigen does not inherit from diffusionpipeline so we hack it
-        self._internal_dict = { # pylint: disable=protected-access
-            'vae': self.vae,
-            'model': self.model,
-            'processor': self.processor,
-        }
-
-    @classmethod
-    def from_pretrained(cls, model_name, vae_path: str=None, cache_dir: str=None):
-        if not os.path.exists(model_name):
-            cache_dir = cache_dir or os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_dir,
-                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
-        model = OmniGen.from_pretrained(model_name)
-        processor = OmniGenProcessor.from_pretrained(model_name)
-        if os.path.exists(os.path.join(model_name, "vae")):
-            vae = AutoencoderKL.from_pretrained(os.path.join(model_name, "vae"))
-        else:
-            vae = AutoencoderKL.from_pretrained(vae_path or "stabilityai/sdxl-vae")
-        return cls(vae, model, processor)
-
-    def merge_lora(self, lora_path: str):
-        model = PeftModel.from_pretrained(self.model, lora_path)
-        model.merge_and_unload()
-        self.model = model
-
-    def to(self, device: Union[str, torch.device]):
-        if isinstance(device, str):
-            device = torch.device(device)
-        self.model.to(device)
-        self.vae.to(device)
-
-    def vae_encode(self, x, dtype):
-        x = x.to(dtype)
-        if self.vae.config.shift_factor is not None:
-            x = self.vae.encode(x).latent_dist.sample()
-            x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor
-        else:
-            x = self.vae.encode(x).latent_dist.sample().mul_(self.vae.config.scaling_factor)
-        x = x.to(dtype)
-        return x
-
-    def move_to_device(self, data):
-        if isinstance(data, list):
-            return [x.to(self.device) for x in data]
-        return data.to(self.device)
-
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        prompt: Union[str, List[str]],
-        input_images: Union[List[str], List[List[str]]] = None,
-        height: int = 1024,
-        width: int = 1024,
-        num_inference_steps: int = 50,
-        guidance_scale: float = 3,
-        use_img_guidance: bool = True,
-        img_guidance_scale: float = 1.6,
-        output_type: str = 'latent',
-        seed: int = None,
-        ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            prompt (`str` or `List[str]`):
-                The prompt or prompts to guide the image generation. 
-            input_images (`List[str]` or `List[List[str]]`, *optional*):
-                The list of input images. We will replace the "<|image_i|>" in prompt with the 1-th image in list.
-            height (`int`, *optional*, defaults to 1024):
-                The height in pixels of the generated image. The number must be a multiple of 16.
-            width (`int`, *optional*, defaults to 1024):
-                The width in pixels of the generated image. The number must be a multiple of 16.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference.
-            guidance_scale (`float`, *optional*, defaults to 4.0):
-                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
-                `guidance_scale` is defined as `w` of equation 2. of [Imagen
-                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
-                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
-                usually at the expense of lower image quality.
-            use_img_guidance (`bool`, *optional*, defaults to True):
-                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). 
-            img_guidance_scale (`float`, *optional*, defaults to 1.6):
-                Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). 
-            self.separate_cfg_infer (`bool`, *optional*, defaults to False):
-                Perform inference on images with different guidance separately; this can save memory when generating images of large size at the expense of slower inference.
-            self.use_kv_cache (`bool`, *optional*, defaults to True): enable kv cache to speed up the inference
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-        Examples:
-
-        Returns:
-            A list with the generated images.
-        """
-        assert height%16 == 0 and width%16 == 0
-        if self.separate_cfg_infer:
-            self.use_kv_cache = False
-            # raise "Currently, don't support both self.use_kv_cache and self.separate_cfg_infer"
-        if input_images is None:
-            use_img_guidance = False
-        if isinstance(prompt, str):
-            prompt = [prompt]
-            input_images = [input_images] if input_images is not None else None
-
-        input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=self.separate_cfg_infer)
-
-        num_prompt = len(prompt)
-        num_cfg = 2 if use_img_guidance else 1
-        latent_size_h, latent_size_w = height//8, width//8
-
-        if seed is not None:
-            generator = torch.Generator(device=self.device).manual_seed(int(seed))
-        else:
-            generator = None
-        latents = torch.randn(num_prompt, 4, latent_size_h, latent_size_w, device=self.device, generator=generator)
-        latents = torch.cat([latents]*(1+num_cfg), 0).to(self.dtype)
-
-        input_img_latents = []
-        if self.separate_cfg_infer:
-            for temp_pixel_values in input_data['input_pixel_values']:
-                temp_input_latents = []
-                for img in temp_pixel_values:
-                    img = self.vae_encode(img.to(self.device), self.dtype)
-                    temp_input_latents.append(img)
-                input_img_latents.append(temp_input_latents)
-        else:
-            for img in input_data['input_pixel_values']:
-                img = self.vae_encode(img.to(self.device), self.dtype)
-                input_img_latents.append(img)
-
-        model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']),
-            input_img_latents=input_img_latents,
-            input_image_sizes=input_data['input_image_sizes'],
-            attention_mask=self.move_to_device(input_data["attention_mask"]),
-            position_ids=self.move_to_device(input_data["position_ids"]),
-            cfg_scale=guidance_scale,
-            img_cfg_scale=img_guidance_scale,
-            use_img_cfg=use_img_guidance,
-            use_kv_cache=self.use_kv_cache)
-
-        if self.separate_cfg_infer:
-            func = self.model.forward_with_separate_cfg
-        else:
-            func = self.model.forward_with_cfg
-        self.model.to(self.dtype)
-
-        scheduler = OmniGenScheduler(num_steps=num_inference_steps)
-        samples = scheduler(latents, func, model_kwargs, use_kv_cache=self.use_kv_cache)
-        samples = samples.chunk((1+num_cfg), dim=0)[0]
-
-        if output_type == 'latent':
-            output_images = { 'images': samples }
-            return output_images
-
-        samples = samples.to(self.vae.dtype)
-        if self.vae.config.shift_factor is not None:
-            samples = samples / self.vae.config.scaling_factor + self.vae.config.shift_factor
-        else:
-            samples = samples / self.vae.config.scaling_factor
-        samples = self.vae.decode(samples).sample
-
-        output_samples = (samples * 0.5 + 0.5).clamp(0, 1)*255
-        output_samples = output_samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy()
-        output_images = []
-        for _i, sample in enumerate(output_samples):
-            output_images.append(Image.fromarray(sample))
-
-        return output_images

modules/omnigen/processor.py

@@ -1,312 +0,0 @@
-import os
-import re
-from typing import Dict, List
-import torch
-from torchvision import transforms
-from transformers import AutoTokenizer
-from huggingface_hub import snapshot_download
-from .utils import crop_arr
-
-
-class OmniGenProcessor:
-    def __init__(self,
-                text_tokenizer,
-                max_image_size: int=1024):
-        self.text_tokenizer = text_tokenizer
-        self.max_image_size = max_image_size
-
-        self.image_transform = transforms.Compose([
-            transforms.Lambda(lambda pil_image: crop_arr(pil_image, max_image_size)),
-            transforms.ToTensor(),
-            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True)
-        ])
-
-        self.collator = OmniGenCollator()
-        self.separate_collator = OmniGenSeparateCollator()
-
-    @classmethod
-    def from_pretrained(cls, model_name):
-        if not os.path.exists(model_name):
-            cache_folder = os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_folder,
-                                           allow_patterns="*.json")
-        text_tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-        return cls(text_tokenizer)
-
-
-    def process_image(self, image):
-        return self.image_transform(image)
-
-    def process_multi_modal_prompt(self, text, input_images):
-        text = self.add_prefix_instruction(text)
-        if input_images is None or len(input_images) == 0:
-            model_inputs = self.text_tokenizer(text)
-            return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None}
-
-        pattern = r"<\|image_\d+\|>"
-        prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)]
-
-        for i in range(1, len(prompt_chunks)):
-            if prompt_chunks[i][0] == 1:
-                prompt_chunks[i] = prompt_chunks[i][1:]
-
-        image_tags = re.findall(pattern, text)
-        image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags]
-
-        unique_image_ids = sorted(list(set(image_ids)))
-        assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}"
-        # total images must be the same as the number of image tags
-        assert len(unique_image_ids) == len(input_images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images"
-
-        input_images = [input_images[x-1] for x in image_ids]
-
-        all_input_ids = []
-        img_inx = []
-        _idx = 0
-        for i in range(len(prompt_chunks)):
-            all_input_ids.extend(prompt_chunks[i])
-            if i != len(prompt_chunks) -1:
-                start_inx = len(all_input_ids)
-                size = input_images[i].size(-2) *  input_images[i].size(-1) // 16 // 16
-                img_inx.append([start_inx, start_inx+size])
-                all_input_ids.extend([0]*size)
-
-        return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx}
-
-
-    def add_prefix_instruction(self, prompt):
-        user_prompt = '<|user|>\n'
-        generation_prompt = 'Generate an image according to the following instructions\n'
-        assistant_prompt = '<|assistant|>\n<|diffusion|>'
-        prompt_suffix = "<|end|>\n"
-        prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}"
-        return prompt
-
-
-    def __call__(self,
-                instructions: List[str],
-                input_images: List[List[str]] = None,
-                height: int = 1024,
-                width: int = 1024,
-                negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.",
-                use_img_cfg: bool = True,
-                separate_cfg_input: bool = False,
-                ) -> Dict:
-
-        if input_images is None:
-            use_img_cfg = False
-        if isinstance(instructions, str):
-            instructions = [instructions]
-            input_images = [input_images]
-
-        input_data = []
-        for i in range(len(instructions)):
-            cur_instruction = instructions[i]
-            cur_input_images = None if input_images is None else input_images[i]
-            if cur_input_images is not None and len(cur_input_images) > 0:
-                cur_input_images = [self.process_image(x) for x in cur_input_images]
-            else:
-                cur_input_images = None
-                assert "<img><|image_1|></img>" not in cur_instruction
-
-            mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images)
-
-            neg_mllm_input, img_cfg_mllm_input = None, None
-            neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None)
-            if use_img_cfg:
-                if cur_input_images is not None and len(cur_input_images) >= 1:
-                    img_cfg_prompt = [f"<img><|image_{i+1}|></img>" for i in range(len(cur_input_images))]
-                    img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images)
-                else:
-                    img_cfg_mllm_input = neg_mllm_input
-
-            input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width]))
-
-        if separate_cfg_input:
-            return self.separate_collator(input_data)
-        return self.collator(input_data)
-
-
-class OmniGenCollator:
-    def __init__(self, pad_token_id=2, hidden_size=3072):
-        self.pad_token_id = pad_token_id
-        self.hidden_size = hidden_size
-
-    def create_position(self, attention_mask, num_tokens_for_output_images):
-        position_ids = []
-        text_length = attention_mask.size(-1)
-        img_length = max(num_tokens_for_output_images)
-        for mask in attention_mask:
-            temp_l = torch.sum(mask)
-            temp_position = [0]*(text_length-temp_l) + [i for i in range(temp_l+img_length+1)] # we add a time embedding into the sequence, so add one more token
-            position_ids.append(temp_position)
-        return torch.LongTensor(position_ids)
-
-    def create_mask(self, attention_mask, num_tokens_for_output_images):
-        extended_mask = []
-        padding_images = []
-        text_length = attention_mask.size(-1)
-        img_length = max(num_tokens_for_output_images)
-        seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token
-        inx = 0
-        for mask in attention_mask:
-            temp_l = torch.sum(mask)
-            pad_l = text_length - temp_l
-
-            temp_mask = torch.tril(torch.ones(size=(temp_l+1, temp_l+1)))
-
-            image_mask = torch.zeros(size=(temp_l+1, img_length))
-            temp_mask = torch.cat([temp_mask, image_mask], dim=-1)
-
-            image_mask = torch.ones(size=(img_length, temp_l+img_length+1))
-            temp_mask = torch.cat([temp_mask, image_mask], dim=0)
-
-            if pad_l > 0:
-                pad_mask = torch.zeros(size=(temp_l+1+img_length, pad_l))
-                temp_mask = torch.cat([pad_mask, temp_mask], dim=-1)
-
-                pad_mask = torch.ones(size=(pad_l, seq_len))
-                temp_mask = torch.cat([pad_mask, temp_mask], dim=0)
-
-            true_img_length = num_tokens_for_output_images[inx]
-            pad_img_length = img_length - true_img_length
-            if pad_img_length > 0:
-                temp_mask[:, -pad_img_length:] = 0
-                temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size))
-            else:
-                temp_padding_imgs = None
-
-            extended_mask.append(temp_mask.unsqueeze(0))
-            padding_images.append(temp_padding_imgs)
-            inx += 1
-        return torch.cat(extended_mask, dim=0), padding_images
-
-    def adjust_attention_for_input_images(self, attention_mask, image_sizes):
-        for b_inx in image_sizes.keys():
-            for start_inx, end_inx in image_sizes[b_inx]:
-                attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1
-
-        return attention_mask
-
-    def pad_input_ids(self, input_ids, image_sizes):
-        max_l = max([len(x) for x in input_ids])
-        padded_ids = []
-        attention_mask = []
-        _new_image_sizes = []
-
-        for i in range(len(input_ids)):
-            temp_ids = input_ids[i]
-            temp_l = len(temp_ids)
-            pad_l = max_l - temp_l
-            if pad_l == 0:
-                attention_mask.append([1]*max_l)
-                padded_ids.append(temp_ids)
-            else:
-                attention_mask.append([0]*pad_l+[1]*temp_l)
-                padded_ids.append([self.pad_token_id]*pad_l+temp_ids)
-
-            if i in image_sizes:
-                new_inx = []
-                for old_inx in image_sizes[i]:
-                    new_inx.append([x+pad_l for x in old_inx])
-                image_sizes[i] = new_inx
-
-        return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes
-
-
-    def process_mllm_input(self, mllm_inputs, target_img_size):
-        num_tokens_for_output_images = []
-        for img_size in target_img_size:
-            num_tokens_for_output_images.append(img_size[0]*img_size[1]//16//16)
-
-        pixel_values, image_sizes = [], {}
-        b_inx = 0
-        for x in mllm_inputs:
-            if x['pixel_values'] is not None:
-                pixel_values.extend(x['pixel_values'])
-                for size in x['image_sizes']:
-                    if b_inx not in image_sizes:
-                        image_sizes[b_inx] = [size]
-                    else:
-                        image_sizes[b_inx].append(size)
-            b_inx += 1
-        pixel_values = [x.unsqueeze(0) for x in pixel_values]
-
-        input_ids = [x['input_ids'] for x in mllm_inputs]
-        padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes)
-        position_ids = self.create_position(attention_mask, num_tokens_for_output_images)
-        attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images)
-        attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes)
-
-        return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes
-
-    def __call__(self, features):
-        mllm_inputs = [f[0] for f in features]
-        cfg_mllm_inputs = [f[1] for f in features]
-        img_cfg_mllm_input = [f[2] for f in features]
-        target_img_size = [f[3] for f in features]
-
-        if img_cfg_mllm_input[0] is not None:
-            mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input
-            target_img_size = target_img_size + target_img_size + target_img_size
-        else:
-            mllm_inputs = mllm_inputs + cfg_mllm_inputs
-            target_img_size = target_img_size + target_img_size
-
-
-        all_padded_input_ids, all_position_ids, all_attention_mask, all_padding_images, all_pixel_values, all_image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
-
-        data = {"input_ids": all_padded_input_ids,
-        "attention_mask": all_attention_mask,
-        "position_ids": all_position_ids,
-        "input_pixel_values": all_pixel_values,
-        "input_image_sizes": all_image_sizes,
-        "padding_images": all_padding_images,
-        }
-        return data
-
-
-class OmniGenSeparateCollator(OmniGenCollator):
-    def __call__(self, features):
-        mllm_inputs = [f[0] for f in features]
-        cfg_mllm_inputs = [f[1] for f in features]
-        img_cfg_mllm_input = [f[2] for f in features]
-        target_img_size = [f[3] for f in features]
-
-        all_padded_input_ids, all_attention_mask, all_position_ids, all_pixel_values, all_image_sizes, all_padding_images = [], [], [], [], [], []
-
-        padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(mllm_inputs, target_img_size)
-        all_padded_input_ids.append(padded_input_ids)
-        all_attention_mask.append(attention_mask)
-        all_position_ids.append(position_ids)
-        all_pixel_values.append(pixel_values)
-        all_image_sizes.append(image_sizes)
-        all_padding_images.append(padding_images)
-
-        if cfg_mllm_inputs[0] is not None:
-            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(cfg_mllm_inputs, target_img_size)
-            all_padded_input_ids.append(padded_input_ids)
-            all_attention_mask.append(attention_mask)
-            all_position_ids.append(position_ids)
-            all_pixel_values.append(pixel_values)
-            all_image_sizes.append(image_sizes)
-            all_padding_images.append(padding_images)
-        if img_cfg_mllm_input[0] is not None:
-            padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(img_cfg_mllm_input, target_img_size)
-            all_padded_input_ids.append(padded_input_ids)
-            all_attention_mask.append(attention_mask)
-            all_position_ids.append(position_ids)
-            all_pixel_values.append(pixel_values)
-            all_image_sizes.append(image_sizes)
-            all_padding_images.append(padding_images)
-
-        data = {"input_ids": all_padded_input_ids,
-        "attention_mask": all_attention_mask,
-        "position_ids": all_position_ids,
-        "input_pixel_values": all_pixel_values,
-        "input_image_sizes": all_image_sizes,
-        "padding_images": all_padding_images,
-        }
-        return data

modules/omnigen/scheduler.py

@@ -1,55 +0,0 @@
-import torch
-from tqdm import tqdm
-from transformers.cache_utils import Cache, DynamicCache
-
-class OmniGenScheduler:
-    def __init__(self, num_steps: int=50, time_shifting_factor: int=1):
-        self.num_steps = num_steps
-        self.time_shift = time_shifting_factor
-
-        t = torch.linspace(0, 1, num_steps+1)
-        t = t / (t + time_shifting_factor - time_shifting_factor * t)
-        self.sigma = t
-    
-    def crop_kv_cache(self, past_key_values, num_tokens_for_img):
-        crop_past_key_values = ()
-        for layer_idx in range(len(past_key_values)):
-            key_states, value_states = past_key_values[layer_idx][:2]
-            crop_past_key_values += ((key_states[..., :-(num_tokens_for_img+1), :], value_states[..., :-(num_tokens_for_img+1), :], ),)
-        return crop_past_key_values
-        # return DynamicCache.from_legacy_cache(crop_past_key_values)
-
-    def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img):
-        if isinstance(position_ids, list):
-            for i in range(len(position_ids)):
-                position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):]
-        else:
-            position_ids = position_ids[:, -(num_tokens_for_img+1):]
-        return position_ids
-
-    def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img):
-        if isinstance(attention_mask, list):
-            return [x[..., -(num_tokens_for_img+1):, :] for x in attention_mask]
-        return attention_mask[..., -(num_tokens_for_img+1):, :]
-
-    def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True):
-        past_key_values = None
-        for i in tqdm(range(self.num_steps)):
-            timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i]
-            pred, temp_past_key_values = func(z, timesteps, past_key_values=past_key_values, **model_kwargs)
-            sigma_next = self.sigma[i+1]
-            sigma = self.sigma[i]
-            z = z + (sigma_next - sigma) * pred
-            if i == 0 and use_kv_cache:
-                num_tokens_for_img = z.size(-1)*z.size(-2) // 4
-                if isinstance(temp_past_key_values, list):
-                    past_key_values = [self.crop_kv_cache(x, num_tokens_for_img) for x in temp_past_key_values]
-                    model_kwargs['input_ids'] = [None] * len(temp_past_key_values)
-                else:
-                    past_key_values = self.crop_kv_cache(temp_past_key_values, num_tokens_for_img)
-                    model_kwargs['input_ids'] = None
-
-                model_kwargs['position_ids'] = self.crop_position_ids_for_cache(model_kwargs['position_ids'], num_tokens_for_img)
-                model_kwargs['attention_mask'] = self.crop_attention_mask_for_cache(model_kwargs['attention_mask'], num_tokens_for_img)
-        return z
-

modules/omnigen/transformer.py

@@ -1,159 +0,0 @@
-import math
-import warnings
-from typing import List, Optional, Tuple, Union
-
-import torch
-import torch.utils.checkpoint
-from torch import nn
-from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-from huggingface_hub import snapshot_download
-
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-    SequenceClassifierOutputWithPast,
-    TokenClassifierOutput,
-)
-from transformers.modeling_utils import PreTrainedModel
-from transformers import Phi3Config, Phi3Model
-from transformers.cache_utils import Cache, DynamicCache, StaticCache
-from transformers.utils import logging
-
-logger = logging.get_logger(__name__)
-
-
-class Phi3Transformer(Phi3Model):
-    """
-    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`]
-    We only modified the attention mask
-    Args:
-        config: Phi3Config
-    """
-
-    def forward(
-        self,
-        input_ids: torch.LongTensor = None,
-        attention_mask: Optional[torch.Tensor] = None,
-        position_ids: Optional[torch.LongTensor] = None,
-        past_key_values: Optional[List[torch.FloatTensor]] = None,
-        inputs_embeds: Optional[torch.FloatTensor] = None,
-        use_cache: Optional[bool] = None,
-        output_attentions: Optional[bool] = None,
-        output_hidden_states: Optional[bool] = None,
-        return_dict: Optional[bool] = None,
-        cache_position: Optional[torch.LongTensor] = None,
-    ) -> Union[Tuple, BaseModelOutputWithPast]:
-        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
-        output_hidden_states = (
-            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
-        )
-        use_cache = use_cache if use_cache is not None else self.config.use_cache
-
-        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
-
-        if (input_ids is None) ^ (inputs_embeds is not None):
-            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
-
-        if self.gradient_checkpointing and self.training:
-            if use_cache:
-                logger.warning_once(
-                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
-                )
-                use_cache = False
-
-        # kept for BC (non `Cache` `past_key_values` inputs)
-        return_legacy_cache = False
-        if use_cache and not isinstance(past_key_values, Cache):
-            return_legacy_cache = True
-            if past_key_values is None:
-                past_key_values = DynamicCache()
-            else:
-                past_key_values = DynamicCache.from_legacy_cache(past_key_values)
-                logger.warning_once(
-                    "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and "
-                    "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class "
-                    "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)"
-                )
-
-        if inputs_embeds is None:
-            inputs_embeds = self.embed_tokens(input_ids)
-
-        if cache_position is None:
-            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
-            cache_position = torch.arange(
-                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
-            )
-        if position_ids is None:
-            position_ids = cache_position.unsqueeze(0)
-
-        if attention_mask is not None and attention_mask.dim() == 3:
-            dtype = inputs_embeds.dtype
-            min_dtype = torch.finfo(dtype).min
-            attention_mask = (1 - attention_mask) * min_dtype
-            attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype)
-        else:
-            raise 
-            # causal_mask = self._update_causal_mask(
-            #     attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
-            # )
-
-        hidden_states = inputs_embeds
-
-        # decoder layers
-        all_hidden_states = () if output_hidden_states else None
-        all_self_attns = () if output_attentions else None
-        next_decoder_cache = None
-
-        for decoder_layer in self.layers:
-            if output_hidden_states:
-                all_hidden_states += (hidden_states,)
-
-            if self.gradient_checkpointing and self.training:
-                layer_outputs = self._gradient_checkpointing_func(
-                    decoder_layer.__call__,
-                    hidden_states,
-                    attention_mask,
-                    position_ids,
-                    past_key_values,
-                    output_attentions,
-                    use_cache,
-                    cache_position,
-                )
-            else:
-                layer_outputs = decoder_layer(
-                    hidden_states,
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    past_key_value=past_key_values,
-                    output_attentions=output_attentions,
-                    use_cache=use_cache,
-                    cache_position=cache_position,
-                )
-
-            hidden_states = layer_outputs[0]
-
-            if use_cache:
-                next_decoder_cache = layer_outputs[2 if output_attentions else 1]
-
-            if output_attentions:
-                all_self_attns += (layer_outputs[1],)
-
-        hidden_states = self.norm(hidden_states)
-
-        # add hidden states from the last decoder layer
-        if output_hidden_states:
-            all_hidden_states += (hidden_states,)
-
-        next_cache = next_decoder_cache if use_cache else None
-        if return_legacy_cache:
-            next_cache = next_cache.to_legacy_cache()
-
-        if not return_dict:
-            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
-        return BaseModelOutputWithPast(
-            last_hidden_state=hidden_states,
-            past_key_values=next_cache,
-            hidden_states=all_hidden_states,
-            attentions=all_self_attns,
-        )
-

modules/omnigen/utils.py

@@ -1,105 +0,0 @@
-import logging
-
-from PIL import Image
-import torch
-import numpy as np
-
-def create_logger(logging_dir):
-    """
-    Create a logger that writes to a log file and stdout.
-    """
-    logging.basicConfig(
-        level=logging.INFO,
-        format='[\033[34m%(asctime)s\033[0m] %(message)s',
-        datefmt='%Y-%m-%d %H:%M:%S',
-        handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")]
-    )
-    logger = logging.getLogger(__name__)
-    return logger
-
-
-@torch.no_grad()
-def update_ema(ema_model, model, decay=0.9999):
-    """
-    Step the EMA model towards the current model.
-    """
-    ema_params = dict(ema_model.named_parameters())
-    for name, param in model.named_parameters():
-        ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay)
-
-
-def requires_grad(model, flag=True):
-    """
-    Set requires_grad flag for all parameters in a model.
-    """
-    for p in model.parameters():
-        p.requires_grad = flag
-
-
-def center_crop_arr(pil_image, image_size):
-    """
-    Center cropping implementation from ADM.
-    https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126
-    """
-    while min(*pil_image.size) >= 2 * image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    scale = image_size / min(*pil_image.size)
-    pil_image = pil_image.resize(
-        tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-    )
-
-    arr = np.array(pil_image)
-    crop_y = (arr.shape[0] - image_size) // 2
-    crop_x = (arr.shape[1] - image_size) // 2
-    return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size])
-
-
-def crop_arr(pil_image, max_image_size):
-    while min(*pil_image.size) >= 2 * max_image_size:
-        pil_image = pil_image.resize(
-            tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    if max(*pil_image.size) > max_image_size:
-        scale = max_image_size / max(*pil_image.size)
-        pil_image = pil_image.resize(
-            tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    if min(*pil_image.size) < 16:
-        scale = 16 / min(*pil_image.size)
-        pil_image = pil_image.resize(
-            tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS
-        )
-
-    arr = np.array(pil_image)
-    crop_y1 = (arr.shape[0] % 16) // 2
-    crop_y2 = arr.shape[0] % 16 - crop_y1
-
-    crop_x1 = (arr.shape[1] % 16) // 2
-    crop_x2 = arr.shape[1] % 16 - crop_x1
-
-    arr = arr[crop_y1:arr.shape[0]-crop_y2, crop_x1:arr.shape[1]-crop_x2]
-    return Image.fromarray(arr)
-
-
-def vae_encode(vae, x, weight_dtype):
-    if x is not None:
-        if vae.config.shift_factor is not None:
-            x = vae.encode(x).latent_dist.sample()
-            x = (x - vae.config.shift_factor) * vae.config.scaling_factor
-        else:
-            x = vae.encode(x).latent_dist.sample().mul_(vae.config.scaling_factor)
-        x = x.to(weight_dtype)
-    return x
-
-
-def vae_encode_list(vae, x, weight_dtype):
-    latents = []
-    for img in x:
-        img = vae_encode(vae, img, weight_dtype)
-        latents.append(img)
-    return latents

modules/paths.py

@@ -29,6 +29,7 @@ script_path = os.path.dirname(modules_path)
 data_path = cli.data_dir
 models_config = cli.models_dir or config.get('models_dir') or 'models'
 models_path = models_config if os.path.isabs(models_config) else os.path.join(data_path, models_config)
+params_path = os.environ.get('SD_PATH_PARAMS', os.path.join(data_path, "params.txt"))
 extensions_dir = cli.extensions_dir or os.path.join(data_path, "extensions")
 extensions_builtin_dir = "extensions-builtin"
 sd_configs_path = os.path.join(script_path, "configs")

modules/postprocess/gfpgan_model.py

@@ -72,8 +72,8 @@ def setup_model(dirname):
     except Exception:
         pass
     try:
-        install('basicsr', quiet=True)
-        install('gfpgan', quiet=True)
+        install('git+https://github.com/Disty0/BasicSR@2b6a12c28e0c81bfb13b7e984144f0b0f5461484', 'basicsr')
+        install('git+https://github.com/Disty0/GFPGAN@09b1190eabbc77e5f15c61fa7c38a2064b403e20', 'gfpgan')
         import gfpgan
         import facexlib
         import modules.detailer

modules/processing_args.py

@@ -166,7 +166,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t
     extra_networks.activate(p, include=['text_encoder', 'text_encoder_2', 'text_encoder_3'])
     if 'prompt' in possible:
         if 'OmniGen' in model.__class__.__name__:
-            prompts = [p.replace('|image|', '<|image_1|>') for p in prompts]
+            prompts = [p.replace('|image|', '<img><|image_1|></img>') for p in prompts]
         if 'HiDreamImage' in model.__class__.__name__  and prompt_parser_diffusers.embedder is not None:
             args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds')
             prompt_embeds = prompt_parser_diffusers.embedder('prompt_embeds')

modules/processing_diffusers.py

@@ -170,7 +170,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
         prev_job = shared.state.job
 
         # hires runs on original pipeline
-        if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None:
+        if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and not shared.opts.control_hires:
             shared.sd_model.restore_pipeline()
 
         # upscale
@@ -200,8 +200,8 @@ def process_hires(p: processing.StableDiffusionProcessing, output):
             if 'Upscale' in shared.sd_model.__class__.__name__ or 'Flux' in shared.sd_model.__class__.__name__ or 'Kandinsky' in shared.sd_model.__class__.__name__:
                 output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height)
             if p.is_control and hasattr(p, 'task_args') and p.task_args.get('image', None) is not None:
-                if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
-                    output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input
+                    if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:
+                        output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input
             update_sampler(p, shared.sd_model, second_pass=True)
             orig_denoise = p.denoising_strength
             p.denoising_strength = strength

modules/processing_helpers.py

@@ -572,10 +572,7 @@ def update_sampler(p, sd_model, second_pass=False):
     if hasattr(sd_model, 'scheduler'):
         if sampler_selection == 'None':
             return
-        if sampler_selection is None:
-            sampler = sd_samplers.all_samplers_map.get("UniPC")
-        else:
-            sampler = sd_samplers.all_samplers_map.get(sampler_selection, None)
+        sampler = sd_samplers.find_sampler(sampler_selection)
         if sampler is None:
             shared.log.warning(f'Sampler: sampler="{sampler_selection}" not found')
             sampler = sd_samplers.all_samplers_map.get("UniPC")

modules/processing_vae.py

@@ -117,6 +117,7 @@ def full_vae_decode(latents, model):
     elif shared.opts.diffusers_offload_mode != "sequential":
         sd_models.move_model(model.vae, devices.device)
 
+    sd_models.set_vae_options(model, vae=None, op='decode')
     upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae)
     if upcast:
         if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent'
@@ -193,6 +194,7 @@ def full_vae_encode(image, model):
     vae_name = sd_vae.loaded_vae_file if sd_vae.loaded_vae_file is not None else "default"
     log_debug(f'Encode vae="{vae_name}" dtype={model.vae.dtype} upcast={model.vae.config.get("force_upcast", None)}')
 
+    sd_models.set_vae_options(model, vae=None, op='encode')
     upcast = (model.vae.dtype == torch.float16) and (getattr(model.vae.config, 'force_upcast', False) or shared.opts.no_half_vae)
     if upcast:
         if hasattr(model, 'upcast_vae'): # this is done by diffusers automatically if output_type != 'latent'
@@ -254,8 +256,8 @@ def vae_postprocess(tensor, model, output_type='np'):
                 if output_type == "pil":
                     images = model.numpy_to_pil(images)
             else:
-                import diffusers
-                model.image_processor = diffusers.image_processor.VaeImageProcessor()
+                from diffusers.image_processor import VaeImageProcessor
+                model.image_processor = VaeImageProcessor()
                 images = model.image_processor.postprocess(tensor, output_type=output_type)
         else:
             images = tensor if isinstance(tensor, list) or isinstance(tensor, np.ndarray) else [tensor]

modules/rocm.py

@@ -81,11 +81,9 @@ class Agent:
         self.blaslt_supported = os.path.exists(os.path.join(blaslt_tensile_libpath, f"Kernels.so-000-{name}.hsaco" if sys.platform == "win32" else f"extop_{name}.co"))
 
     def get_gfx_version(self) -> Union[str, None]:
-        if self.gfx_version >= 0x1200:
-            return "12.0.0"
-        elif self.gfx_version >= 0x1100:
+        if self.gfx_version >= 0x1101 and self.gfx_version < 0x1200:
             return "11.0.0"
-        elif self.gfx_version >= 0x1000:
+        elif self.gfx_version != 0x1030 and self.gfx_version >= 0x1000 and self.gfx_version < 0x1100:
             # gfx1010 users had to override gfx version to 10.3.0 in Linux
             # it is unknown whether overriding is needed in ZLUDA
             return "10.3.0"
@@ -206,7 +204,7 @@ else:
         if agent.gfx_version >= 0x1100 and os.environ.get("FLASH_ATTENTION_USE_TRITON_ROCM", "false").lower() != "true":
             # use the navi_rotary_fix fork because the original doesn't support rotary_emb for transformers
             # original: "git+https://github.com/ROCm/flash-attention@howiejay/navi_support"
-            default = "https://github.com/Disty0/flash-attention@navi_rotary_fix"
+            default = "git+https://github.com/Disty0/flash-attention@navi_rotary_fix"
         return os.environ.get("FLASH_ATTENTION_PACKAGE", default)
 
     is_wsl: bool = os.environ.get('WSL_DISTRO_NAME', 'unknown' if spawn('wslpath -w /') else None) is not None

modules/sd_hijack_accelerate.py

@@ -36,7 +36,7 @@ def hijack_set_module_tensor(
         # note: majority of time is spent on .to(old_value.dtype)
         if tensor_name in module._buffers: # pylint: disable=protected-access
             module._buffers[tensor_name] = value.to(device, old_value.dtype)  # pylint: disable=protected-access
-        elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device):  # pylint: disable=protected-access
+        elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device):  # pylint: disable=protected-access
             param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
             module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access
     t1 = time.time()
@@ -64,7 +64,7 @@ def hijack_set_module_tensor_simple(
     with devices.inference_context():
         if tensor_name in module._buffers: # pylint: disable=protected-access
             module._buffers[tensor_name] = value.to(device)  # pylint: disable=protected-access
-        elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device):  # pylint: disable=protected-access
+        elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device):  # pylint: disable=protected-access
             param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access
             module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access
     t1 = time.time()

modules/sd_models.py

@@ -92,7 +92,7 @@ def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False):
         if shared.opts.diffusers_vae_upcast != 'default':
             sd_model.vae.config.force_upcast = True if shared.opts.diffusers_vae_upcast == 'true' else False
             shared.log.quiet(quiet, f'Setting {op}: component=VAE upcast={sd_model.vae.config.force_upcast}')
-        if shared.opts.no_half_vae:
+        if shared.opts.no_half_vae and op not in {'decode', 'encode'}:
             devices.dtype_vae = torch.float32
             sd_model.vae.to(devices.dtype_vae)
             shared.log.quiet(quiet, f'Setting {op}: component=VAE no-half=True')
@@ -105,11 +105,20 @@ def set_vae_options(sd_model, vae=None, op:str='model', quiet:bool=False):
     if hasattr(sd_model, "enable_vae_tiling"):
         if shared.opts.diffusers_vae_tiling:
             if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'config') and hasattr(sd_model.vae.config, 'sample_size') and isinstance(sd_model.vae.config.sample_size, int):
+                if getattr(sd_model.vae, "tile_sample_min_size_backup", None) is None:
+                    sd_model.vae.tile_sample_min_size_backup = sd_model.vae.tile_sample_min_size
+                    sd_model.vae.tile_latent_min_size_backup = sd_model.vae.tile_latent_min_size
+                    sd_model.vae.tile_overlap_factor_backup = sd_model.vae.tile_overlap_factor
                 if shared.opts.diffusers_vae_tile_size > 0:
                     sd_model.vae.tile_sample_min_size = int(shared.opts.diffusers_vae_tile_size)
-                    sd_model.vae.tile_latent_min_size = int(sd_model.vae.config.sample_size / (2 ** (len(sd_model.vae.config.block_out_channels) - 1)))
+                    sd_model.vae.tile_latent_min_size = int(shared.opts.diffusers_vae_tile_size / (2 ** (len(sd_model.vae.config.block_out_channels) - 1)))
+                else:
+                    sd_model.vae.tile_sample_min_size = getattr(sd_model.vae, "tile_sample_min_size_backup", sd_model.vae.tile_sample_min_size)
+                    sd_model.vae.tile_latent_min_size = getattr(sd_model.vae, "tile_latent_min_size_backup", sd_model.vae.tile_latent_min_size)
                 if shared.opts.diffusers_vae_tile_overlap != 0.25:
                     sd_model.vae.tile_overlap_factor = float(shared.opts.diffusers_vae_tile_overlap)
+                else:
+                    sd_model.vae.tile_overlap_factor = getattr(sd_model.vae, "tile_overlap_factor_backup", sd_model.vae.tile_overlap_factor)
                 shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True tile={sd_model.vae.tile_sample_min_size} overlap={sd_model.vae.tile_overlap_factor}')
             else:
                 shared.log.quiet(quiet, f'Setting {op}: component=VAE tiling=True')
@@ -582,7 +591,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         if "Kandinsky" in sd_model.__class__.__name__: # need a special case
             sd_model.scheduler.name = 'DDIM'
 
-        if model_type not in ['Stable Cascade']: # need a special-case
+        if hasattr(sd_model, "unet") and model_type not in ['Stable Cascade']: # others calls load_diffuser again
             sd_unet.load_unet(sd_model)
 
         add_noise_pred_to_diffusers_callback(sd_model)
@@ -651,7 +660,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
         from modules import modelstats
         modelstats.analyze()
 
-    shared.log.info(f"Load {op}: time={timer.summary()} native={get_native(sd_model)} memory={memory_stats()}")
+    shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.dct()} native={get_native(sd_model)} memory={memory_stats()}")
 
 
 class DiffusersTaskType(Enum):
@@ -899,32 +908,45 @@ def set_diffuser_pipe(pipe, new_pipe_type):
 def set_diffusers_attention(pipe, quiet:bool=False):
     import diffusers.models.attention_processor as p
 
-    def set_attn(pipe, attention):
+    def set_attn(pipe, attention, name:str=None, quiet:bool=False):
         if attention is None:
             return
         # other models uses their own attention processor
         if pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet"):
             pipe.unet.set_attn_processor(attention)
+        elif not quiet:
+            shared.log.warning(f"Attention: {name if name is not None else attention.__class__.__name__} is not compatible with {pipe.__class__.__name__}")
 
     # if hasattr(pipe, 'pipe'):
     #    set_diffusers_attention(pipe.pipe)
 
-    if 'ControlNet' in pipe.__class__.__name__: # do not replace attention in ControlNet pipelines
+    if 'Control' in pipe.__class__.__name__ or 'Adapter' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")):
+        if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}:
+            shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}")
+        else:
+            pipe.current_attn_name = shared.opts.cross_attention_optimization
         return
+
     shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
     if shared.opts.cross_attention_optimization == "Disabled":
         pass # do nothing
     elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
-        set_attn(pipe, p.AttnProcessor2_0())
-    elif shared.opts.cross_attention_optimization == "xFormers" and hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
-        pipe.enable_xformers_memory_efficient_attention()
-    elif shared.opts.cross_attention_optimization == "Split attention" and hasattr(pipe, "enable_attention_slicing"):
-        pipe.enable_attention_slicing()
+        set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product", quiet=True)
+    elif shared.opts.cross_attention_optimization == "xFormers":
+        if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
+            pipe.enable_xformers_memory_efficient_attention()
+        else:
+            shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
+    elif shared.opts.cross_attention_optimization == "Split attention":
+        if hasattr(pipe, "enable_attention_slicing"):
+            pipe.enable_attention_slicing()
+        else:
+            shared.log.warning(f"Attention: Split attention is not compatible with {pipe.__class__.__name__}")
     elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
-        set_attn(pipe, p.AttnProcessor())
+        set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix")
     elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
         from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
-        set_attn(pipe, DynamicAttnProcessorBMM())
+        set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM")
 
     pipe.current_attn_name = shared.opts.cross_attention_optimization
 
@@ -1064,7 +1086,8 @@ def clear_caches():
         lora_common.loaded_networks.clear()
         lora_common.previously_loaded_networks.clear()
         lora_load.lora_cache.clear()
-    from modules import prompt_parser_diffusers, memstats
+    from modules import prompt_parser_diffusers, memstats, sd_offload
+    sd_offload.offload_hook_instance = None
     prompt_parser_diffusers.cache.clear()
     memstats.reset_stats()
 

modules/sd_models_utils.py

@@ -155,6 +155,8 @@ def apply_function_to_model(sd_model, function, options, op=None):
         if hasattr(sd_model, 'transformer') and hasattr(sd_model.transformer, 'config'):
             sd_model.transformer = function(sd_model.transformer, op="transformer", sd_model=sd_model)
     if "Model" in options:
+        if hasattr(sd_model, 'model') and (hasattr(sd_model.model, 'config') or isinstance(sd_model.model, torch.nn.Module)):
+            sd_model.model = function(sd_model.model, op="model", sd_model=sd_model)
         if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
             sd_model.unet = function(sd_model.unet, op="unet", sd_model=sd_model)
         if hasattr(sd_model, 'decoder_pipe') and hasattr(sd_model, 'decoder'):

modules/sd_offload.py

@@ -4,6 +4,7 @@ import time
 import inspect
 import torch
 import accelerate.hooks
+import accelerate.utils.modeling
 from installer import log
 from modules import shared, devices, errors, model_quant
 from modules.timer import process as process_timer
@@ -14,7 +15,17 @@ debug_move = log.trace if debug else lambda *args, **kwargs: None
 offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4']
 offload_post = ['h1']
 offload_hook_instance = None
-balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline']
+balanced_offload_exclude = ['CogView4Pipeline']
+accelerate_dtype_byte_size = None
+
+
+def dtype_byte_size(dtype: torch.dtype):
+    try:
+        if dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz]:
+            dtype = accelerate.utils.modeling.CustomDtype.FP8
+    except Exception: # catch since older torch many not have defined dtypes
+        pass
+    return accelerate_dtype_byte_size(dtype)
 
 
 def get_signature(cls):
@@ -58,6 +69,7 @@ def set_accelerate(sd_model):
 
 
 def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
+    global accelerate_dtype_byte_size # pylint: disable=global-statement
     t0 = time.time()
     if not shared.native:
         shared.log.warning('Attempting to use offload with backend=original')
@@ -67,6 +79,9 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False):
         return
     if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate):
         sd_model.has_accelerate = False
+    if accelerate_dtype_byte_size is None:
+        accelerate_dtype_byte_size = accelerate.utils.modeling.dtype_byte_size
+        accelerate.utils.modeling.dtype_byte_size = dtype_byte_size
     if shared.opts.diffusers_offload_mode == "none":
         if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type:
             shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model')
@@ -156,21 +171,25 @@ class OffloadHook(accelerate.hooks.ModelHook):
         return module
 
     def pre_forward(self, module, *args, **kwargs):
-        if devices.normalize_device(module.device) != devices.normalize_device(devices.device):
+        if not devices.same_device(module.device, devices.device):
             device_index = torch.device(devices.device).index
             if device_index is None:
                 device_index = 0
             max_memory = { device_index: self.gpu, "cpu": self.cpu }
             device_map = getattr(module, "balanced_offload_device_map", None)
             if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None):
+                # try:
                 device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory)
+                # except Exception as e:
+                #     shared.log.error(f'Offload: type=balanced module={module.__class__.__name__} {e}')
             offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
             if devices.backend == "directml":
                 keys = device_map.keys()
                 for v in keys:
                     if isinstance(device_map[v], int):
                         device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type
-            module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
+            if device_map is not None:
+                module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
             module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             module.balanced_offload_device_map = device_map
             module.balanced_offload_max_memory = max_memory

modules/sd_samplers.py

@@ -16,6 +16,17 @@ flow_models =  ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogVi
 flow_models += ['Hunyuan', 'LTX', 'Mochi']
 
 
+def find_sampler(name:str):
+    if name is None or name == 'None':
+        return all_samplers_map.get("UniPC", None)
+    for sampler in all_samplers:
+        if sampler.name.lower() == name.lower() or name in sampler.aliases:
+            debug(f'Find sampler: name="{name}" found={sampler.name}')
+            return sampler
+    debug(f'Find sampler: name="{name}" found=None')
+    return None
+
+
 def list_samplers():
     global all_samplers # pylint: disable=global-statement
     global all_samplers_map # pylint: disable=global-statement

modules/sd_unet.py

@@ -34,7 +34,7 @@ def load_unet(model):
             if prior_text_encoder is not None:
                 model.prior_pipe.text_encoder = None # Prevent OOM
                 model.prior_pipe.text_encoder = prior_text_encoder.to(devices.device, dtype=devices.dtype)
-        elif "Flux" in model.__class__.__name__ or "StableDiffusion3" in model.__class__.__name__ or "HiDream" in model.__class__.__name__:
+        elif "Flux" in model.__class__.__name__ or "StableDiffusion3" in model.__class__.__name__ or "HiDream" in model.__class__.__name__ or "Lumina2" in model.__class__.__name__:
             loaded_unet = shared.opts.sd_unet
             sd_models.load_diffuser() # TODO model load: force-reloading entire model as loading transformers only leads to massive memory usage
             """
@@ -71,7 +71,7 @@ def load_unet(model):
 
 def refresh_unet_list():
     unet_dict.clear()
-    for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors", ".gguf"]):
+    for file in files_cache.list_files(shared.opts.unet_dir, ext_filter=[".safetensors", ".gguf", ".pth"]):
         basename = os.path.basename(file)
         name = os.path.splitext(basename)[0] if ".safetensors" in basename else basename
         unet_dict[name] = file

modules/sd_vae_remote.py

@@ -12,14 +12,25 @@ hf_decode_endpoints = {
     'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud',
     'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud',
     'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
-    'h1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud',
     'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud',
 }
+hf_decode_endpoints['pixartalpha'] = hf_decode_endpoints['sd']
+hf_decode_endpoints['pixartsigma'] = hf_decode_endpoints['sdxl']
+hf_decode_endpoints['omnigen'] = hf_decode_endpoints['sdxl']
+hf_decode_endpoints['h1'] = hf_decode_endpoints['f1']
+hf_decode_endpoints['lumina2'] = hf_decode_endpoints['f1']
+
 hf_encode_endpoints = {
     'sd': 'https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud',
     'sdxl': 'https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud',
     'f1': 'https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud',
 }
+hf_encode_endpoints['pixartalpha'] = hf_encode_endpoints['sd']
+hf_encode_endpoints['pixartsigma'] = hf_encode_endpoints['sdxl']
+hf_encode_endpoints['omnigen'] = hf_encode_endpoints['sdxl']
+hf_encode_endpoints['h1'] = hf_encode_endpoints['f1']
+hf_encode_endpoints['lumina2'] = hf_encode_endpoints['f1']
+
 dtypes = {
     "float16": torch.float16,
     "float32": torch.float32,
@@ -74,7 +85,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_
                 params["output_type"] = "pt"
                 params["output_tensor_type"] = "binary"
                 headers["Accept"] = "tensor/binary"
-            if (model_type == 'f1' or model_type == 'h1') and (width > 0) and (height > 0):
+            if model_type in {'f1', 'h1', 'lumina2'} and (width > 0) and (height > 0):
                 params['width'] = width
                 params['height'] = height
             if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None:

modules/sd_vae_taesd.py

@@ -36,7 +36,7 @@ prev_cls = ''
 prev_type = ''
 prev_model = ''
 lock = threading.Lock()
-supported = ['sd', 'sdxl', 'f1', 'h1', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha']
+supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha', 'omnigen']
 
 
 def warn_once(msg, variant=None):
@@ -52,36 +52,36 @@ def warn_once(msg, variant=None):
 def get_model(model_type = 'decoder', variant = None):
     global prev_cls, prev_type, prev_model # pylint: disable=global-statement
     from modules import shared
-    cls = shared.sd_model_type
-    if cls == 'ldm': # original backend
-        cls = 'sd'
-    elif cls == 'h1': # hidream uses flux vae
-        cls = 'f1'
-    elif cls == 'pixartsigma':
-        cls = 'sdxl'
-    elif cls == 'pixartalpha':
-        cls = 'sd'
-    elif cls not in supported:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+    model_cls = shared.sd_model_type
+    if model_cls is None or model_cls == 'none':
+        return None
+    elif model_cls in {'ldm', 'pixartalpha'}:
+        model_cls = 'sd'
+    elif model_cls in {'h1', 'lumina2'}:
+        model_cls = 'f1'
+    elif model_cls in {'pixartsigma', 'omnigen'}:
+        model_cls = 'sdxl'
+    elif model_cls not in supported:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     variant = variant or shared.opts.taesd_variant
     folder = os.path.join(paths.models_path, "TAESD")
     os.makedirs(folder, exist_ok=True)
     if variant.startswith('TAE'):
         cfg = TAESD_MODELS[variant]
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
-        fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth')
+        fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth')
         if not os.path.exists(fn):
             uri = cfg['uri']
             if not uri.endswith('.pth'):
-                uri += '/tae' + cls + '_' + model_type + '.pth'
+                uri += '/tae' + model_cls + '_' + model_type + '.pth'
             try:
                 shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download')
                 torch.hub.download_url_to_file(uri, fn)
             except Exception as e:
                 warn_once(f'download uri={uri} {e}', variant=variant)
         if os.path.exists(fn):
-            prev_cls = cls
+            prev_cls = model_cls
             prev_type = model_type
             prev_model = variant
             shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load')
@@ -99,14 +99,14 @@ def get_model(model_type = 'decoder', variant = None):
                 TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
             return TAESD_MODELS[variant]['model']
     elif variant.startswith('Hybrid'):
-        cfg = CQYAN_MODELS[variant].get(cls, None)
-        if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+        cfg = CQYAN_MODELS[variant].get(model_cls, None)
+        if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
             return cfg['model']
         if cfg is None:
-            warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+            warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
             return None
         repo = cfg['repo']
-        prev_cls = cls
+        prev_cls = model_cls
         prev_type = model_type
         prev_model = variant
         shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load')
@@ -118,10 +118,12 @@ def get_model(model_type = 'decoder', variant = None):
             from modules.taesd.hybrid_small import AutoencoderSmall
             vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype)
         vae = vae.to(devices.device, dtype=dtype)
-        CQYAN_MODELS[variant][cls]['model'] = vae
+        CQYAN_MODELS[variant][model_cls]['model'] = vae
         return vae
+    elif variant is None:
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} variant is none', variant=variant)
     else:
-        warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant)
+        warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
     return None
 
 

modules/sdnq/__init__.py

@@ -0,0 +1,445 @@
+# 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 torch
+from diffusers.quantizers.base import DiffusersQuantizer
+from diffusers.quantizers.quantization_config import QuantizationConfigMixin
+from diffusers.utils import get_module_from_name
+from modules import devices, shared
+
+from .common import dtype_dict, use_tensorwise_fp8_matmul, quantized_matmul_dtypes, allowed_types, conv_types, conv_transpose_types
+from .dequantizer import dequantizer_dict
+from .forward import get_forward_func
+
+
+def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None):
+    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
+        result_shape = None
+        original_shape = layer.weight.shape
+        if torch_dtype is None:
+            torch_dtype = devices.dtype
+
+        if layer_class_name in conv_types:
+            if not quant_conv:
+                return layer
+            if dtype_dict[weights_dtype]["num_bits"] < 4:
+                weights_dtype = "uint4"
+            is_conv_type = True
+            reduction_axes = 1
+            output_channel_size, channel_size = layer.weight.shape[:2]
+            group_channel_size = channel_size // layer.groups
+            use_quantized_matmul = False
+            if use_quantized_matmul_conv:
+                use_quantized_matmul = weights_dtype in quantized_matmul_dtypes and group_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 group_channel_size % 16 == 0
+                if use_quantized_matmul:
+                    result_shape = layer.weight.shape
+                    layer.weight.data = layer.weight.reshape(output_channel_size, -1)
+        elif layer_class_name in conv_transpose_types:
+            if not quant_conv:
+                return layer
+            if dtype_dict[weights_dtype]["num_bits"] < 4:
+                weights_dtype = "uint4"
+            is_conv_transpose_type = True
+            reduction_axes = 0
+            channel_size, output_channel_size = layer.weight.shape[:2]
+            use_quantized_matmul = False
+        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:
+                    if dtype_dict[weights_dtype]["is_integer"]:
+                        use_quantized_matmul = output_channel_size % 8 == 0 and channel_size % 8 == 0
+                    else:
+                        use_quantized_matmul = output_channel_size % 16 == 0 and channel_size % 16 == 0
+
+        if group_size == 0:
+            if is_linear_type:
+                group_size = 2 ** (2 + dtype_dict[weights_dtype]["num_bits"])
+            else:
+                group_size = 2 ** (1 + dtype_dict[weights_dtype]["num_bits"])
+        elif group_size != -1 and not is_linear_type:
+            group_size = max(group_size // 2, 1)
+
+        if not use_quantized_matmul and group_size > 0:
+            if group_size >= channel_size:
+                group_size = channel_size
+                num_of_groups = 1
+            else:
+                num_of_groups = channel_size // group_size
+                while num_of_groups * group_size != channel_size: # 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
+            group_size = int(group_size)
+            num_of_groups = int(num_of_groups)
+
+            if num_of_groups > 1:
+                result_shape = layer.weight.shape
+                new_shape = list(result_shape)
+                if is_conv_type:
+                    # output_channel_size, channel_size, X, X
+                    # output_channel_size, num_of_groups, group_size, X, X
+                    new_shape[1] = group_size
+                    new_shape.insert(1, num_of_groups)
+                    reduction_axes = 2
+                elif is_conv_transpose_type:
+                    #channel_size, output_channel_size, X, X
+                    #num_of_groups, group_size, output_channel_size, X, X
+                    new_shape[0] = group_size
+                    new_shape.insert(0, num_of_groups)
+                    reduction_axes = 1
+                elif is_linear_type:
+                    # output_channel_size, channel_size
+                    # output_channel_size, num_of_groups, group_size
+                    last_dim_index = layer.weight.ndim
+                    new_shape[last_dim_index - 1 : last_dim_index] = (num_of_groups, group_size)
+                layer.weight.data = layer.weight.reshape(new_shape)
+
+        layer.weight.requires_grad = False
+        if return_device is None:
+            return_device = layer.weight.device
+        if quantization_device is not None:
+            layer.weight.data = layer.weight.to(quantization_device)
+        if layer.weight.dtype != torch.float32:
+            layer.weight.data = layer.weight.to(dtype=torch.float32)
+
+        layer.weight.data, scale, zero_point = quantize_weight(layer.weight, reduction_axes, weights_dtype)
+        if not dequantize_fp32 and not (use_quantized_matmul and not dtype_dict[weights_dtype]["is_integer"] and not use_tensorwise_fp8_matmul):
+            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_dequantizer = dequantizer_dict[weights_dtype](
+            scale=scale,
+            zero_point=zero_point,
+            quantized_weight_shape=layer.weight.shape,
+            result_dtype=torch_dtype,
+            result_shape=result_shape,
+            original_shape=original_shape,
+            weights_dtype=weights_dtype,
+            use_quantized_matmul=use_quantized_matmul,
+        )
+        layer.weight.data = layer.sdnq_dequantizer.pack_weight(layer.weight).to(return_device)
+        layer.sdnq_dequantizer = layer.sdnq_dequantizer.to(return_device)
+
+        layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul)
+        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, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None, modules_to_not_convert: List[str] = []): # pylint: disable=unused-argument
+    has_children = list(model.children())
+    if not has_children:
+        return model
+    for module_param_name, module in model.named_children():
+        if module_param_name in modules_to_not_convert:
+            continue
+        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,
+                use_quantized_matmul_conv=use_quantized_matmul_conv,
+                dequantize_fp32=dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
+                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,
+                use_quantized_matmul_conv=use_quantized_matmul_conv,
+                dequantize_fp32=dequantize_fp32,
+                quantization_device=quantization_device,
+                return_device=return_device,
+                param_name=module_param_name,
+                modules_to_not_convert=modules_to_not_convert,
+            )
+    return model
+
+
+def get_scale_asymmetric(weight: torch.FloatTensor, reduction_axes: Union[int, 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: Union[int, List[int]], weights_dtype: str) -> torch.FloatTensor:
+    scale = torch.amax(weight.abs(), dim=reduction_axes, keepdims=True).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, reduction_axes: Union[int, List[int]], weights_dtype: str) -> Tuple[torch.Tensor, torch.FloatTensor, torch.FloatTensor]:
+    if dtype_dict[weights_dtype]["is_unsigned"]:
+        scale, zero_point = get_scale_asymmetric(weight, reduction_axes, weights_dtype)
+        quantized_weight = torch.sub(weight, zero_point).div_(scale)
+    else:
+        scale = get_scale_symmetric(weight, reduction_axes, weights_dtype)
+        quantized_weight = torch.div(weight, scale)
+        zero_point = None
+    if dtype_dict[weights_dtype]["is_integer"]:
+        quantized_weight.round_()
+    quantized_weight = quantized_weight.clamp_(dtype_dict[weights_dtype]["min"], dtype_dict[weights_dtype]["max"]).to(dtype_dict[weights_dtype]["torch_dtype"])
+    return quantized_weight, scale, zero_point
+
+
+class QuantizationMethod(str, Enum):
+    SDNQ = "sdnq"
+
+
+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):
+        super().__init__(quantization_config, **kwargs)
+        self.modules_to_not_convert = []
+
+    def check_if_quantized_param(
+        self,
+        model,
+        param_value: "torch.Tensor",
+        param_name: str,
+        state_dict: Dict[str, Any], # pylint: disable=unused-argument
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        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, # pylint: disable=unused-argument
+    ):
+        if self.quantization_config.return_device is not None:
+            return_device = self.quantization_config.return_device
+        else:
+            return_device = target_device
+
+        if self.quantization_config.quantization_device is not None:
+            target_device = self.quantization_config.quantization_device
+
+        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, _ = get_module_from_name(model, param_name)
+        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,
+            use_quantized_matmul_conv=self.quantization_config.use_quantized_matmul_conv,
+            dequantize_fp32=self.quantization_config.dequantize_fp32,
+            quantization_device=None,
+            return_device=return_device,
+            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.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, # pylint: disable=unused-argument
+    ):
+        if keep_in_fp32_modules is not None:
+            self.modules_to_not_convert.extend(keep_in_fp32_modules)
+        self.modules_to_not_convert.extend(self.quantization_config.modules_to_not_convert)
+        self.quantization_config.modules_to_not_convert = self.modules_to_not_convert
+        model.config.quantization_config = self.quantization_config
+
+    def _process_model_after_weight_loading(self, model, **kwargs): # pylint: disable=unused-argument
+        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 True
+
+    @property
+    def is_compileable(self):
+        return True
+
+
+@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", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "float8_e4m3fn", "float8_e4m3fnuz", "float8_e5m2", "float8_e5m2fnuz")
+        weights_dtype (`int`, *optional*, defaults to `0`):
+            Used to decide how many elements of a tensor will share the same quantization group.
+        quant_conv (`bool`, *optional*, defaults to `False`):
+            Enabling this option will quantize the convolutional layers in UNet models too.
+        use_quantized_matmul (`bool`, *optional*, defaults to `False`):
+            Enabling this option will use quantized INT8 or FP8 MatMul instead of BF16 / FP16.
+        use_quantized_matmul_conv (`bool`, *optional*, defaults to `False`):
+            Same as use_quantized_matmul_conv but for the convolutional layers with UNets like SDXL.
+        dequantize_fp32 (`bool`, *optional*, defaults to `False`):
+            Enabling this option will use FP32 on the dequantization step.
+        quantization_device (`torch.device`, *optional*, defaults to `None`):
+            Used to set which device will be used for the quantization calculation on model load.
+        return_device (`torch.device`, *optional*, defaults to `None`):
+            Used to set which device will the quantized weights be sent back to.
+        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,
+        use_quantized_matmul_conv: bool = False,
+        dequantize_fp32: bool = False,
+        quantization_device: Optional[torch.device] = None,
+        return_device: Optional[torch.device] = None,
+        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.use_quantized_matmul_conv = use_quantized_matmul_conv
+        self.dequantize_fp32 = dequantize_fp32
+        self.quantization_device = quantization_device
+        self.return_device = return_device
+        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", "int7", "int6", "int5", "int4", "int3", "int2", "uint8", "uint7", "uint6", "uint5", "uint4", "uint3", "uint2", "uint1", "bool", "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}")
+        if not isinstance(self.modules_to_not_convert, list):
+            self.modules_to_not_convert = [self.modules_to_not_convert]

modules/sdnq/common.py

@@ -0,0 +1,42 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+import sys
+import torch
+from modules import devices
+
+torch_version = float(torch.__version__[:3])
+
+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},
+    "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True},
+    "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "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},
+    "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True},
+    "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": "uint2", "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},
+}
+dtype_dict["bool"] = dtype_dict["uint1"]
+if hasattr(torch, "float8_e4m3fnuz"):
+    dtype_dict["float8_e4m3fnuz"] = {"min": -240, "max": 240, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}
+if hasattr(torch, "float8_e5m2fnuz"):
+    dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}
+
+use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9))
+quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "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

modules/sdnq/dequantizer.py

@@ -0,0 +1,193 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+import torch
+from modules import shared
+
+from .common import dtype_dict
+from .packed_int import pack_int_symetric, unpack_int_symetric, packed_int_function_dict
+
+
+def dequantize_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    result = torch.addcmul(zero_point, weight.to(dtype=scale.dtype), scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def dequantize_symmetric(weight: torch.CharTensor, scale: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+    if skip_quantized_matmul:
+        result = weight.transpose(0,1).to(dtype=scale.dtype).mul_(scale.transpose(0,1)).to(dtype=dtype)
+    else:
+        result = weight.to(dtype=scale.dtype).mul_(scale).to(dtype=dtype)
+    if result_shape is not None:
+        result = result.reshape(result_shape)
+    return result
+
+
+def dequantize_symmetric_with_bias(weight: torch.CharTensor, scale: torch.FloatTensor, bias: torch.FloatTensor, dtype: torch.dtype, result_shape: torch.Size) -> torch.FloatTensor:
+    return torch.addcmul(bias, weight.to(dtype=scale.dtype), scale).to(dtype=dtype).reshape(result_shape)
+
+
+def dequantize_packed_int_asymmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, zero_point: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str) -> torch.FloatTensor:
+    return dequantize_asymmetric(packed_int_function_dict[weights_dtype]["unpack"](weight, shape), scale, zero_point, dtype, result_shape)
+
+
+def dequantize_packed_int_symmetric(weight: torch.ByteTensor, scale: torch.FloatTensor, shape: torch.Size, dtype: torch.dtype, result_shape: torch.Size, weights_dtype: str, skip_quantized_matmul: bool = False) -> torch.FloatTensor:
+    if skip_quantized_matmul:
+        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale.transpose(0,1), dtype, result_shape)
+    else:
+        return dequantize_symmetric(unpack_int_symetric(weight, shape, weights_dtype, dtype=scale.dtype), scale, dtype, result_shape)
+
+
+class AsymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.FloatTensor,
+        zero_point: torch.FloatTensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        original_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
+        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): # pylint: disable=unused-argument
+        return dequantize_asymmetric_compiled(weight, self.scale, self.zero_point, self.result_dtype, self.result_shape)
+
+
+class SymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.FloatTensor,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        original_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
+        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): # pylint: disable=unused-argument
+        return dequantize_symmetric_compiled(weight, self.scale, self.result_dtype, self.result_shape, skip_quantized_matmul=skip_quantized_matmul)
+
+
+class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.FloatTensor,
+        zero_point: torch.FloatTensor,
+        quantized_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        original_shape: torch.Size,
+        weights_dtype: str,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.use_quantized_matmul = False
+        self.original_shape = original_shape
+        self.quantized_weight_shape = quantized_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): # pylint: disable=unused-argument
+        return dequantize_packed_int_asymmetric_compiled(weight, self.scale, self.zero_point, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype)
+
+
+class PackedINTSymmetricWeightsDequantizer(torch.nn.Module):
+    def __init__(
+        self,
+        scale: torch.FloatTensor,
+        quantized_weight_shape: torch.Size,
+        result_dtype: torch.dtype,
+        result_shape: torch.Size,
+        original_shape: torch.Size,
+        weights_dtype: str,
+        use_quantized_matmul: bool = False,
+        **kwargs, # pylint: disable=unused-argument
+    ):
+        super().__init__()
+        self.weights_dtype = weights_dtype
+        self.original_shape = original_shape
+        self.use_quantized_matmul = use_quantized_matmul
+        self.quantized_weight_shape = quantized_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): # pylint: disable=unused-argument
+        return dequantize_packed_int_symmetric_compiled(weight, self.scale, self.quantized_weight_shape, self.result_dtype, self.result_shape, self.weights_dtype, skip_quantized_matmul=skip_quantized_matmul)
+
+
+dequantizer_dict = {
+    "int8": SymmetricWeightsDequantizer,
+    "int7": PackedINTSymmetricWeightsDequantizer,
+    "int6": PackedINTSymmetricWeightsDequantizer,
+    "int5": PackedINTSymmetricWeightsDequantizer,
+    "int4": PackedINTSymmetricWeightsDequantizer,
+    "int3": PackedINTSymmetricWeightsDequantizer,
+    "int2": PackedINTSymmetricWeightsDequantizer,
+    "uint8": AsymmetricWeightsDequantizer,
+    "uint7": PackedINTAsymmetricWeightsDequantizer,
+    "uint6": PackedINTAsymmetricWeightsDequantizer,
+    "uint5": PackedINTAsymmetricWeightsDequantizer,
+    "uint4": PackedINTAsymmetricWeightsDequantizer,
+    "uint3": PackedINTAsymmetricWeightsDequantizer,
+    "uint2": PackedINTAsymmetricWeightsDequantizer,
+    "uint1": AsymmetricWeightsDequantizer,
+    "bool": AsymmetricWeightsDequantizer,
+    "float8_e4m3fn": SymmetricWeightsDequantizer,
+    "float8_e4m3fnuz": SymmetricWeightsDequantizer,
+    "float8_e5m2": SymmetricWeightsDequantizer,
+    "float8_e5m2fnuz": SymmetricWeightsDequantizer,
+}
+
+
+if shared.opts.sdnq_dequantize_compile:
+    try:
+        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
+        dequantize_asymmetric_compiled = torch.compile(dequantize_asymmetric, fullgraph=True)
+        dequantize_symmetric_compiled = torch.compile(dequantize_symmetric, fullgraph=True)
+        dequantize_packed_int_asymmetric_compiled = torch.compile(dequantize_packed_int_asymmetric, fullgraph=True)
+        dequantize_packed_int_symmetric_compiled = torch.compile(dequantize_packed_int_symmetric, fullgraph=True)
+    except Exception as e:
+        shared.log.warning(f"Quantization: type=sdnq Dequantize using torch.compile is not available: {e}")
+        dequantize_asymmetric_compiled = dequantize_asymmetric
+        dequantize_symmetric_compiled = dequantize_symmetric
+        dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric
+        dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric
+else:
+    dequantize_asymmetric_compiled = dequantize_asymmetric
+    dequantize_symmetric_compiled = dequantize_symmetric
+    dequantize_packed_int_asymmetric_compiled = dequantize_packed_int_asymmetric
+    dequantize_packed_int_symmetric_compiled = dequantize_packed_int_symmetric

modules/sdnq/forward.py

@@ -0,0 +1,409 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Callable, List, Tuple, Optional
+import torch
+from modules import shared
+
+from .common import conv_types, conv_transpose_types
+from .dequantizer import dequantize_symmetric, dequantize_symmetric_with_bias
+from .packed_int import unpack_int_symetric
+
+
+def get_forward_func(layer_class_name: str, use_quantized_matmul: bool, is_integer: bool, use_tensorwise_fp8_matmul: bool) -> Callable: # pylint: disable=inconsistent-return-statements
+    if layer_class_name in conv_types:
+        if use_quantized_matmul:
+            if is_integer:
+                return quantized_conv_forward_int8_matmul
+            else:
+                if use_tensorwise_fp8_matmul:
+                    return quantized_conv_forward_fp8_matmul_tensorwise
+                else:
+                    return quantized_conv_forward_fp8_matmul
+        else:
+            return quantized_conv_forward
+    elif layer_class_name in conv_transpose_types:
+        if layer_class_name.endswith("1d"):
+            return quantized_conv_transpose_1d_forward
+        elif layer_class_name.endswith("2d"):
+            return quantized_conv_transpose_2d_forward
+        elif layer_class_name.endswith("3d"):
+            return quantized_conv_transpose_3d_forward
+    else:
+        if use_quantized_matmul:
+            if is_integer:
+                return quantized_linear_forward_int8_matmul
+            else:
+                if use_tensorwise_fp8_matmul:
+                    return quantized_linear_forward_fp8_matmul_tensorwise
+                else:
+                    return quantized_linear_forward_fp8_matmul
+        else:
+            return quantized_linear_forward
+
+
+def quantize_fp8_matmul_input(input: torch.FloatTensor) -> Tuple[torch.Tensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(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.Tensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(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.CharTensor, torch.FloatTensor]:
+    input = input.flatten(0,-2).contiguous()
+    input_scale = torch.amax(input.abs(), dim=-1, keepdims=True).div_(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)
+    if bias is not None:
+        return dequantize_symmetric_with_bias(torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype), scale, bias, return_dtype, output_shape)
+    else:
+        return dequantize_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)
+
+
+def int8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    quantized_weight_shape: torch.Size,
+    weights_dtype: str,
+) -> torch.FloatTensor:
+    if quantized_weight_shape is not None:
+        weight = unpack_int_symetric(weight, quantized_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)
+    if bias is not None:
+        return dequantize_symmetric_with_bias(torch._int_mm(input, weight), scale, bias, return_dtype, output_shape)
+    else:
+        return dequantize_symmetric(torch._int_mm(input, weight), scale, return_dtype, output_shape)
+
+
+def process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation):
+    if conv_type == 1:
+        batch_size, _, L_in = input.shape
+        C_out, _, K_l = result_shape
+        L_out = (L_in + 2 * padding[1] - dilation[1] * (K_l - 1) - 1) // stride[1] + 1
+        mm_output_shape = (batch_size, L_out, C_out)
+        kernel_size = (1, K_l)
+    if conv_type == 2:
+        batch_size, _, H_in, W_in = input.shape
+        C_out, _, K_h, K_w = result_shape
+        H_out = (H_in + 2 * padding[0] - dilation[0] * (K_h - 1) - 1) // stride[0] + 1
+        W_out = (W_in + 2 * padding[1] - dilation[1] * (K_w - 1) - 1) // stride[1] + 1
+        mm_output_shape = (batch_size, H_out, W_out, C_out)
+        kernel_size = (K_h, K_w)
+    else:
+        batch_size, _, D_in, H_in, W_in = input.shape
+        C_out, _, K_d, K_h, K_w = result_shape
+        D_out = (D_in + 2 * padding[0] - dilation[0] * (K_d - 1) - 1) // stride[0] + 1
+        H_out = (H_in + 2 * padding[1] - dilation[1] * (K_h - 1) - 1) // stride[1] + 1
+        W_out = (W_in + 2 * padding[2] - dilation[2] * (K_w - 1) - 1) // stride[2] + 1
+        mm_output_shape = (batch_size, D_out, H_out, W_out, C_out)
+        kernel_size = (K_d, K_h, K_w)
+
+    if padding_mode != "zeros":
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice, mode=padding_mode)
+        padding = (0,) * (conv_type if conv_type != 1 else 2)
+    elif conv_type == 3:
+        input = torch.nn.functional.pad(input, reversed_padding_repeated_twice)
+
+    if conv_type == 1:
+        input = input.unsqueeze(2)
+
+    if conv_type == 3:
+        K_D_eff = K_d + (K_d - 1) * (dilation[0] - 1)
+        K_H_eff = K_h + (K_h - 1) * (dilation[0] - 1)
+        K_W_eff = K_w + (K_w - 1) * (dilation[0] - 1)
+        input = input.unfold(2, K_D_eff, stride[0]).unfold(3, K_H_eff, stride[1]).unfold(4, K_W_eff, stride[2])
+        if dilation[0] > 1:
+            input = input[..., ::dilation[0], :, :]
+        if dilation[1] > 1:
+            input = input[..., ::dilation[1], :]
+        if dilation[2] > 1:
+            input = input[..., ::dilation[2]]
+        input = input.permute(0, 2, 3, 4, 1, 5, 6, 7).reshape(batch_size, D_out * H_out * W_out, -1)
+    else:
+        input = torch.nn.functional.unfold(input, kernel_size=kernel_size, padding=padding, stride=stride, dilation=dilation).transpose(1,2)
+    return input, mm_output_shape
+
+
+def conv_fp8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, input_scale = quantize_fp8_matmul_input(input)
+
+    if groups == 1:
+        result = torch._scaled_mm(input, weight, scale_a=input_scale, scale_b=scale, bias=bias, out_dtype=return_dtype).reshape(mm_output_shape)
+    else:
+        scale = scale.reshape(groups, 1, scale.shape[1] // groups)
+        input_scale = input_scale.reshape(groups, input_scale.shape[0] // groups, 1)
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        if bias is not None:
+            bias = bias.reshape(groups, bias.shape[0] // groups)
+            for i in range(groups):
+                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=bias[i], out_dtype=return_dtype))
+        else:
+            for i in range(groups):
+                result.append(torch._scaled_mm(input[i], weight[i], scale_a=input_scale[i], scale_b=scale[i], bias=None, out_dtype=return_dtype))
+        result = torch.cat(result, dim=-1).reshape(mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    return result
+
+
+def conv_fp8_matmul_tensorwise(
+    input: torch.FloatTensor,
+    weight: torch.Tensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, scale = quantize_fp8_matmul_input_tensorwise(input, scale)
+    dummy_input_scale = torch.ones(1, device=input.device, dtype=torch.float32)
+
+    if groups == 1:
+        result = torch._scaled_mm(input, weight, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype)
+    else:
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        for i in range(groups):
+            result.append(torch._scaled_mm(input[i], weight[i], scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=scale.dtype))
+        result = torch.cat(result, dim=-1)
+    if bias is not None:
+        dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
+    else:
+        dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    return result
+
+
+def conv_int8_matmul(
+    input: torch.FloatTensor,
+    weight: torch.CharTensor,
+    bias: torch.FloatTensor,
+    scale: torch.FloatTensor,
+    result_shape: torch.Size,
+    quantized_weight_shape: torch.Size,
+    weights_dtype: str,
+    reversed_padding_repeated_twice: List[int],
+    padding_mode: str, conv_type: int,
+    groups: int, stride: List[int],
+    padding: List[int], dilation: List[int],
+) -> torch.FloatTensor:
+    return_dtype = input.dtype
+    input, mm_output_shape = process_conv_input(conv_type, input, reversed_padding_repeated_twice, padding_mode, result_shape, stride, padding, dilation)
+    input, scale = quantize_int8_matmul_input(input, scale)
+    if quantized_weight_shape is not None:
+        weight = unpack_int_symetric(weight, quantized_weight_shape, weights_dtype, dtype=torch.int8, transpose=True)
+
+    if groups == 1:
+        result = torch._int_mm(input, weight)
+    else:
+        weight = weight.reshape(weight.shape[0], groups, weight.shape[1] // groups).transpose(0,1)
+        input = input.reshape(input.shape[0], groups, input.shape[1] // groups).transpose(0,1)
+        result = []
+        for i in range(groups):
+            result.append(torch._int_mm(input[i], weight[i]))
+        result = torch.cat(result, dim=-1)
+    if bias is not None:
+        result = dequantize_symmetric_with_bias(result, scale, bias, return_dtype, mm_output_shape)
+    else:
+        result = dequantize_symmetric(result, scale, return_dtype, mm_output_shape)
+
+    if conv_type == 1:
+        result = result.transpose(1,2)
+    elif conv_type == 2:
+        result = result.permute(0,3,1,2)
+    elif conv_type == 3:
+        result = result.permute(0,4,1,2,3)
+    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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return fp8_matmul_tensorwise(input, self.weight, self.bias, self.sdnq_dequantizer.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_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    return int8_matmul(input, self.weight, self.bias, self.sdnq_dequantizer.scale, getattr(self.sdnq_dequantizer, "quantized_weight_shape", None), self.sdnq_dequantizer.weights_dtype)
+
+
+def quantized_linear_forward(self, input: torch.FloatTensor) -> torch.FloatTensor:
+    return torch.nn.functional.linear(input, self.sdnq_dequantizer(self.weight), self.bias)
+
+
+def get_conv_args(input_ndim: int, stride, padding, dilation):
+    if input_ndim == 3:
+        conv_type = 1
+    elif input_ndim == 4:
+        conv_type = 2
+    else:
+        conv_type = 3
+    if isinstance(stride, int):
+        stride = (stride,) * conv_type
+    if isinstance(padding, int):
+        padding = (padding,) * conv_type
+    if isinstance(dilation, int):
+        dilation = (dilation,) * conv_type
+    if conv_type == 1:
+        stride = (1, stride[0])
+        padding = (0, padding[0])
+        dilation = (1, dilation[0])
+    return conv_type, stride, padding, dilation
+
+
+def quantized_conv_forward_fp8_matmul(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_fp8_matmul(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward_fp8_matmul_tensorwise(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_fp8_matmul_tensorwise(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward_int8_matmul(self, input) -> torch.FloatTensor:
+    if torch.numel(input) / input.shape[2] < 32:
+        return self._conv_forward(input, self.sdnq_dequantizer(self.weight, skip_quantized_matmul=True), self.bias)
+    conv_type, stride, padding, dilation = get_conv_args(input.ndim, self.stride, self.padding, self.dilation)
+    return conv_int8_matmul(
+        input, self.weight, self.bias,
+        self.sdnq_dequantizer.scale,
+        self.sdnq_dequantizer.result_shape,
+        getattr(self.sdnq_dequantizer, "quantized_weight_shape", None),
+        self.sdnq_dequantizer.weights_dtype,
+        self._reversed_padding_repeated_twice,
+        self.padding_mode, conv_type,
+        self.groups, stride, padding, dilation,
+    )
+
+
+def quantized_conv_forward(self, input) -> torch.FloatTensor:
+    return self._conv_forward(input, self.sdnq_dequantizer(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_dequantizer(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_dequantizer(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_dequantizer(self.weight), self.bias, self.stride, self.padding, output_padding, self.groups, self.dilation)
+
+
+if shared.opts.sdnq_dequantize_compile:
+    try:
+        torch._dynamo.config.cache_size_limit = max(8192, torch._dynamo.config.cache_size_limit)
+        int8_matmul = torch.compile(int8_matmul, fullgraph=True)
+        fp8_matmul = torch.compile(fp8_matmul, fullgraph=True)
+        fp8_matmul_tensorwise = torch.compile(fp8_matmul_tensorwise, fullgraph=True)
+        conv_int8_matmul = torch.compile(conv_int8_matmul, fullgraph=True)
+        conv_fp8_matmul = torch.compile(conv_fp8_matmul, fullgraph=True)
+        conv_fp8_matmul_tensorwise = torch.compile(conv_fp8_matmul_tensorwise, fullgraph=True)
+    except Exception as e:
+        shared.log.warning(f"Quantization: type=sdnq MatMul using torch.compile is not available: {e}")

modules/sdnq/packed_int.py

@@ -0,0 +1,256 @@
+# pylint: disable=redefined-builtin,no-member,protected-access
+
+from typing import Optional
+import torch
+
+from .common import dtype_dict
+
+
+def pack_int_symetric(tensor: torch.CharTensor, weights_dtype: str) -> torch.ByteTensor:
+    return packed_int_function_dict[weights_dtype]["pack"](tensor.sub_(dtype_dict[weights_dtype]["min"]).to(dtype=dtype_dict[weights_dtype]["storage_dtype"]))
+
+
+def unpack_int_symetric(packed_tensor: torch.ByteTensor, shape: torch.Size, weights_dtype: str, dtype: Optional[torch.dtype] = None, transpose: Optional[bool] = False) -> torch.CharTensor:
+    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_uint7(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 1), 128)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 2), 128)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128)),
+            torch.bitwise_or(packed_tensor[:, 3], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128)),
+            torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128)),
+            torch.bitwise_or(packed_tensor[:, 5], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 6), 128)),
+            torch.bitwise_or(packed_tensor[:, 6], torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 7), 128)),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint6(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    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_uint5(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 5], 5)),
+            torch.bitwise_or(packed_tensor[:, 1], torch.bitwise_left_shift(packed_tensor[:, 6], 5)),
+            torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 7], 5)),
+            torch.bitwise_or(
+                packed_tensor[:, 3],
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 5], 2), 96),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 3), 128),
+                ),
+            ),
+            torch.bitwise_or(
+                packed_tensor[:, 4],
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 2), 96),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 4), 128),
+                ),
+            ),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint4(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    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_uint3(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    packed_tensor = tensor.contiguous().reshape(-1, 8)
+    packed_tensor = torch.stack(
+        (
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 0], torch.bitwise_left_shift(packed_tensor[:, 1], 3)),
+                torch.bitwise_left_shift(packed_tensor[:, 6], 6),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 2], torch.bitwise_left_shift(packed_tensor[:, 3], 3)),
+                torch.bitwise_left_shift(packed_tensor[:, 7], 6),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_or(packed_tensor[:, 4], torch.bitwise_left_shift(packed_tensor[:, 5], 3)),
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 6], 4), 64),
+                    torch.bitwise_and(torch.bitwise_left_shift(packed_tensor[:, 7], 5), 128),
+                )
+            ),
+        ),
+        dim=-1
+    )
+    return packed_tensor
+
+
+def pack_uint2(tensor: torch.ByteTensor) -> torch.ByteTensor:
+    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_uint7(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 127),
+            torch.bitwise_and(packed_tensor[:, 1], 127),
+            torch.bitwise_and(packed_tensor[:, 2], 127),
+            torch.bitwise_and(packed_tensor[:, 3], 127),
+            torch.bitwise_and(packed_tensor[:, 4], 127),
+            torch.bitwise_and(packed_tensor[:, 5], 127),
+            torch.bitwise_and(packed_tensor[:, 6], 127),
+            torch.bitwise_or(
+                torch.bitwise_or(
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 1), 64),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 2), 32),
+                    ),
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 16),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 4), 8),
+                    ),
+                ),
+                torch.bitwise_or(
+                    torch.bitwise_or(
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 5), 4),
+                        torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 5], 6), 2),
+                    ),
+                    torch.bitwise_right_shift(packed_tensor[:, 6], 7),
+                ),
+            )
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint6(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    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_uint5(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 31),
+            torch.bitwise_and(packed_tensor[:, 1], 31),
+            torch.bitwise_and(packed_tensor[:, 2], 31),
+            torch.bitwise_and(packed_tensor[:, 3], 31),
+            torch.bitwise_and(packed_tensor[:, 4], 31),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 0], 5),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 2), 24),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 1], 5),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 2), 24),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 2], 5),
+                torch.bitwise_or(
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 3], 3), 16),
+                    torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 4], 4), 8),
+                ),
+            ),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint4(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    result = torch.stack((torch.bitwise_and(packed_tensor, 15), torch.bitwise_right_shift(packed_tensor, 4)), dim=-1).reshape(shape)
+    return result
+
+
+def unpack_uint3(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    result = torch.stack(
+        (
+            torch.bitwise_and(packed_tensor[:, 0], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 0], 3), 7),
+            torch.bitwise_and(packed_tensor[:, 1], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 1], 3), 7),
+            torch.bitwise_and(packed_tensor[:, 2], 7),
+            torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 3), 7),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 0], 6),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 4), 4),
+            ),
+            torch.bitwise_or(
+                torch.bitwise_right_shift(packed_tensor[:, 1], 6),
+                torch.bitwise_and(torch.bitwise_right_shift(packed_tensor[:, 2], 5), 4),
+            ),
+        ),
+        dim=-1
+    ).reshape(shape)
+    return result
+
+
+def unpack_uint2(packed_tensor: torch.ByteTensor, shape: torch.Size) -> torch.ByteTensor:
+    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
+
+
+packed_int_function_dict = {
+    "int7": {"pack": pack_uint7, "unpack": unpack_uint7},
+    "int6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "int5": {"pack": pack_uint5, "unpack": unpack_uint5},
+    "int4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "int3": {"pack": pack_uint3, "unpack": unpack_uint3},
+    "int2": {"pack": pack_uint2, "unpack": unpack_uint2},
+    "uint7": {"pack": pack_uint7, "unpack": unpack_uint7},
+    "uint6": {"pack": pack_uint6, "unpack": unpack_uint6},
+    "uint5": {"pack": pack_uint5, "unpack": unpack_uint5},
+    "uint4": {"pack": pack_uint4, "unpack": unpack_uint4},
+    "uint3": {"pack": pack_uint3, "unpack": unpack_uint3},
+    "uint2": {"pack": pack_uint2, "unpack": unpack_uint2},
+}

modules/shared.py

@@ -66,6 +66,7 @@ dir_timestamps = {}
 dir_cache = {}
 max_workers = 8
 default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub')
+sdnq_quant_modes = ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"]
 
 
 class Backend(Enum):
@@ -518,14 +519,15 @@ 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_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes, "visible": native}),
+    "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['default'] + sdnq_quant_modes, "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_compile": OptionInfo(devices.has_triton(), "Decompress using torch.compile", gr.Checkbox, {"visible": native}),
+    "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}),
     "sdnq_use_quantized_matmul": OptionInfo(False, "Use Quantized MatMul", gr.Checkbox, {"visible": native}),
     "sdnq_use_quantized_matmul_conv": OptionInfo(False, "Use Quantized MatMul with convolutional layers", gr.Checkbox, {"visible": native}),
     "sdnq_quantize_with_gpu": OptionInfo(True, "Quantize with the GPU", gr.Checkbox, {"visible": native}),
-    "sdnq_decompress_fp32": OptionInfo(False, "Decompress using full precision", gr.Checkbox, {"visible": native}),
+    "sdnq_dequantize_fp32": OptionInfo(False, "Dequantize using full precision", 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),
@@ -893,10 +895,11 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), {
 }))
 
 options_templates.update(options_section(('control', "Control Options"), {
-    "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}),
-    "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options"),
-    "control_move_processor": OptionInfo(False, "Processor move to CPU after use"),
-    "control_unload_processor": OptionInfo(False, "Processor unload after use"),
+    "control_hires": OptionInfo(False, "Use control during hires", gr.Checkbox, {"visible": False}),
+    "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1, "visible": False}),
+    "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options", gr.Textbox, {"visible": False}),
+    "control_move_processor": OptionInfo(False, "Processor move to CPU after use", gr.Checkbox, {"visible": False}),
+    "control_unload_processor": OptionInfo(False, "Processor unload after use", gr.Checkbox, {"visible": False}),
 }))
 
 options_templates.update(options_section(('interrogate', "Interrogate"), {

modules/shared_items.py

@@ -29,7 +29,7 @@ pipelines = {
     'FLEX': getattr(diffusers, 'AutoPipelineForText2Image', None),
     'Sana': getattr(diffusers, 'SanaPipeline', None),
     'Lumina-Next': getattr(diffusers, 'LuminaText2ImgPipeline', None),
-    'Lumina 2': getattr(diffusers, 'Lumina2Text2ImgPipeline', None),
+    'Lumina 2': getattr(diffusers, 'Lumina2Pipeline', None),
     'AuraFlow': getattr(diffusers, 'AuraFlowPipeline', None),
     'Kandinsky 2.1': getattr(diffusers, 'KandinskyCombinedPipeline', None),
     'Kandinsky 2.2': getattr(diffusers, 'KandinskyV22CombinedPipeline', None),
@@ -41,10 +41,10 @@ pipelines = {
     'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None),
     'Amused': getattr(diffusers, 'AmusedPipeline', None),
     'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None),
+    'OmniGenPipeline': getattr(diffusers, 'OmniGenPipeline', None),
 
     # dynamically imported and redefined later
     'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
-    'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
     'InstaFlow': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
     'SegMoE': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser
 }

modules/teacache/__init__.py

@@ -1,11 +1,12 @@
 from .teacache_flux import teacache_flux_forward
 from .teacache_hidream import teacache_hidream_forward
+from .teacache_lumina2 import teacache_lumina2_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', 'HiDream']
+supported_models = ['Flux', 'CogVideoX', 'Mochi', 'LTX', 'HiDream', 'Lumina2']
 
 
 def apply_teacache(p):
@@ -25,5 +26,7 @@ def apply_teacache(p):
     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]
+    if shared.sd_model.__class__.__name__.startswith('Lumina2'):
+        shared.sd_model.transformer.__class__.cache = {}
+        shared.sd_model.transformer.__class__.uncond_seq_len = None
     shared.log.info(f'Transformers cache: type=teacache cls={shared.sd_model.__class__.__name__} thresh={shared.opts.teacache_thresh}')

modules/teacache/teacache_hidream.py

@@ -111,7 +111,8 @@ def teacache_hidream_forward(
             should_calc = True
             self.accumulated_rel_l1_distance = 0
         else:
-            rescale_func = np.poly1d(self.coefficients)
+            coefficients = [-3.13605009e+04, -7.12425503e+02, 4.91363285e+01, 8.26515490e+00, 1.08053901e-01]
+            rescale_func = np.poly1d(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

modules/teacache/teacache_lumina2.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import Any, Dict, Optional, Union, List
+
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+def teacache_lumina2_forward(
+    self,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    encoder_hidden_states: torch.Tensor,
+    encoder_attention_mask: torch.Tensor,
+    attention_kwargs: Optional[Dict[str, Any]] = None,
+    return_dict: bool = True,
+) -> Union[torch.Tensor, Transformer2DModelOutput]:
+    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:
+        scale_lora_layers(self, lora_scale)
+
+    batch_size, _, height, width = hidden_states.shape
+    temb, encoder_hidden_states_processed = self.time_caption_embed(hidden_states, timestep, encoder_hidden_states)
+    (image_patch_embeddings, context_rotary_emb, noise_rotary_emb, joint_rotary_emb,
+     encoder_seq_lengths, seq_lengths) = self.rope_embedder(hidden_states, encoder_attention_mask)
+    image_patch_embeddings = self.x_embedder(image_patch_embeddings)
+    for layer in self.context_refiner:
+        encoder_hidden_states_processed = layer(encoder_hidden_states_processed, encoder_attention_mask, context_rotary_emb)
+    for layer in self.noise_refiner:
+        image_patch_embeddings = layer(image_patch_embeddings, None, noise_rotary_emb, temb)
+
+    max_seq_len = max(seq_lengths)
+    input_to_main_loop = image_patch_embeddings.new_zeros(batch_size, max_seq_len, self.config.hidden_size)
+    for i, (enc_len, seq_len_val) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        input_to_main_loop[i, :enc_len] = encoder_hidden_states_processed[i, :enc_len]
+        input_to_main_loop[i, enc_len:seq_len_val] = image_patch_embeddings[i]
+
+    use_mask = len(set(seq_lengths)) > 1
+    attention_mask_for_main_loop_arg = None
+    if use_mask:
+        mask = input_to_main_loop.new_zeros(batch_size, max_seq_len, dtype=torch.bool)
+        for i, (enc_len, seq_len_val) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+            mask[i, :seq_len_val] = True
+        attention_mask_for_main_loop_arg = mask
+
+    should_calc = True
+    if self.enable_teacache:
+        cache_key = max_seq_len
+        if cache_key not in self.cache:
+            self.cache[cache_key] = {
+                "accumulated_rel_l1_distance": 0.0,
+                "previous_modulated_input": None,
+                "previous_residual": None,
+            }
+
+        current_cache = self.cache[cache_key]
+        modulated_inp, _, _, _ = self.layers[0].norm1(input_to_main_loop, temb)
+
+        if self.cnt == 0 or self.cnt == self.num_steps - 1:
+            should_calc = True
+            current_cache["accumulated_rel_l1_distance"] = 0.0
+        else:
+            if current_cache["previous_modulated_input"] is not None:
+                # teacache v1 coefficients:
+                coefficients = [393.76566581, -603.50993606, 209.10239044, -23.00726601, 0.86377344]
+                # teacache v2 coefficients:
+                #coefficients = [225.7042019806413, -608.8453716535591, 304.1869942338369, 124.21267720116742, -1.4089066892956552]
+                rescale_func = np.poly1d(coefficients)
+                prev_mod_input = current_cache["previous_modulated_input"]
+                prev_mean = prev_mod_input.abs().mean()
+
+                if prev_mean.item() > 1e-9:
+                    rel_l1_change = ((modulated_inp - prev_mod_input).abs().mean() / prev_mean).cpu().item()
+                else:
+                    rel_l1_change = 0.0 if modulated_inp.abs().mean().item() < 1e-9 else float('inf')
+
+                current_cache["accumulated_rel_l1_distance"] += rescale_func(rel_l1_change)
+
+                if current_cache["accumulated_rel_l1_distance"] < self.rel_l1_thresh:
+                    should_calc = False
+                else:
+                    should_calc = True
+                    current_cache["accumulated_rel_l1_distance"] = 0.0
+            else:
+                should_calc = True
+                current_cache["accumulated_rel_l1_distance"] = 0.0
+
+        current_cache["previous_modulated_input"] = modulated_inp.clone()
+
+        if self.uncond_seq_len is None:
+            self.uncond_seq_len = cache_key
+        if cache_key != self.uncond_seq_len:
+            self.cnt += 1
+            if self.cnt >= self.num_steps:
+                self.cnt = 0
+
+    if self.enable_teacache and not should_calc:
+        if max_seq_len in self.cache and "previous_residual" in self.cache[max_seq_len] and self.cache[max_seq_len]["previous_residual"] is not None:
+             processed_hidden_states = input_to_main_loop + self.cache[max_seq_len]["previous_residual"]
+        else:
+             should_calc = True
+             current_processing_states = input_to_main_loop
+             for layer in self.layers:
+                current_processing_states = layer(current_processing_states, attention_mask_for_main_loop_arg, joint_rotary_emb, temb)
+             processed_hidden_states = current_processing_states
+
+
+    if not (self.enable_teacache and not should_calc) :
+        current_processing_states = input_to_main_loop
+        for layer in self.layers:
+            current_processing_states = layer(current_processing_states, attention_mask_for_main_loop_arg, joint_rotary_emb, temb)
+
+        if self.enable_teacache:
+            if max_seq_len in self.cache:
+                 self.cache[max_seq_len]["previous_residual"] = current_processing_states - input_to_main_loop
+            else:
+                 logger.warning(f"TeaCache: Cache key {max_seq_len} not found when trying to save residual.")
+
+        processed_hidden_states = current_processing_states
+
+    output_after_norm = self.norm_out(processed_hidden_states, temb)
+    p = self.config.patch_size
+    final_output_list = []
+    for i, (enc_len, seq_len_val) in enumerate(zip(encoder_seq_lengths, seq_lengths)):
+        image_part = output_after_norm[i][enc_len:seq_len_val]
+        h_p, w_p = height // p, width // p
+        reconstructed_image = image_part.view(h_p, w_p, p, p, self.out_channels) \
+                                        .permute(4, 0, 2, 1, 3) \
+                                        .flatten(3, 4) \
+                                        .flatten(1, 2)
+        final_output_list.append(reconstructed_image)
+
+    final_output_tensor = torch.stack(final_output_list, dim=0)
+
+    if USE_PEFT_BACKEND:
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (final_output_tensor,)
+
+    return Transformer2DModelOutput(sample=final_output_tensor)

modules/timer.py

@@ -44,8 +44,7 @@ class Timer:
     def summary(self, min_time=default_min_time, total=True):
         if self.profile:
             min_time = -1
-        if self.total <= 0:
-            self.total = sum(self.records.values())
+        self.total = sum(self.records.values())
         res = f"total={self.total:.2f} " if total else ''
         additions = [x for x in self.records.items() if x[1] >= min_time]
         additions = sorted(additions, key=lambda x: x[1], reverse=True)
@@ -60,6 +59,8 @@ class Timer:
     def dct(self, min_time=default_min_time):
         if self.profile:
             res = {k: round(v, 4) for k, v in self.records.items()}
+        self.total = sum(self.records.values())
+        self.records['total'] = self.total
         res = {k: round(v, 2) for k, v in self.records.items() if v >= min_time}
         res = {k: v for k, v in sorted(res.items(), key=lambda x: x[1], reverse=True)} # noqa: C416 # pylint: disable=unnecessary-comprehension
         return res

modules/ui_common.py

@@ -54,7 +54,7 @@ def infotext_to_html(text):
     return code
 
 
-def delete_files(js_data, files, _html_info, index):
+def delete_files(js_data, files, all_files, index):
     try:
         data = json.loads(js_data)
     except Exception:
@@ -63,25 +63,26 @@ def delete_files(js_data, files, _html_info, index):
     if index > -1 and shared.opts.save_selected_only and (index >= data['index_of_first_image']):
         files = [files[index]]
         start_index = index
-        filenames = []
-    filenames = []
-    fullfns = []
+    deleted = []
+    all_files = [f.split('/file=')[1] if 'file=' in f else f for f in all_files] if isinstance(all_files, list) else []
     for _image_index, filedata in enumerate(files, start_index):
-        if 'name' in filedata and os.path.isfile(filedata['name']):
-            fullfn = filedata['name']
-            filenames.append(os.path.basename(fullfn))
-            try:
-                os.remove(fullfn)
-                base, _ext = os.path.splitext(fullfn)
-                desc = f'{base}.txt'
-                if os.path.exists(desc):
-                    os.remove(desc)
-                fullfns.append(fullfn)
-                shared.log.info(f"Deleting image: {fullfn}")
-            except Exception as e:
-                shared.log.error(f'Error deleting file: {fullfn} {e}')
-    files = [image for image in files if image['name'] not in fullfns]
-    return files, plaintext_to_html(f"Deleted: {filenames[0] if len(filenames) > 0 else 'none'}")
+        try:
+            fn = filedata['name']
+            if os.path.isfile(fn):
+                deleted.append(fn)
+                os.remove(fn)
+                if fn in all_files:
+                    all_files.remove(fn)
+                shared.log.info(f'Delete: image="{fn}"')
+            base, _ext = os.path.splitext(fn)
+            desc = f'{base}.txt'
+            if os.path.exists(desc):
+                os.remove(desc)
+                shared.log.info(f'Delete: text="{fn}"')
+        except Exception as e:
+            shared.log.error(f'Delete: image="{fn}" {e}')
+    deleted = ', '.join(deleted) if len(deleted) > 0 else 'none'
+    return all_files, plaintext_to_html(f"Deleted: {deleted}")
 
 
 def save_files(js_data, files, html_info, index):
@@ -296,8 +297,8 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe
                     inputs=[generation_info, result_gallery, html_info, html_info],
                     outputs=[download_files, html_log],
                 )
-                delete.click(fn=call_queue.wrap_gradio_call(delete_files),show_progress=False,
-                    _js="(x, y, z, i) => [x, y, z, selected_gallery_index()]",
+                delete.click(fn=call_queue.wrap_gradio_call(delete_files), show_progress=False,
+                    _js="(x, y, i, j) => [x, y, ...selected_gallery_files()]",
                     inputs=[generation_info, result_gallery, html_info, html_info],
                     outputs=[result_gallery, html_log],
                 )

modules/ui_control.py

@@ -467,9 +467,31 @@ def create_ui(_blocks: gr.Blocks=None):
                             if i == 0:
                                 units[-1].enabled = True # enable first unit in group
 
-                with gr.Accordion('Processor settings', open=False, elem_classes=['control-settings']) as _tab_settings:
+                with gr.Accordion('Control settings', open=False, elem_classes=['control-settings']) as _tab_settings:
                     with gr.Group(elem_classes=['processor-group']):
                         settings = []
+                        with gr.Accordion('Global', open=True, elem_classes=['processor-settings']):
+                            control_hires = gr.Checkbox(label="Use control during hires", value=shared.opts.control_hires, elem_id='control_hires')
+                            def set_control_hires(value):
+                                shared.opts.control_active = value
+                            control_hires.change(fn=set_control_hires, inputs=[control_hires], outputs=[])
+                            control_max_units = gr.Slider(label="Maximum units", minimum=1, maximum=10, step=1, value=shared.opts.control_max_units, elem_id='control_max_units')
+                            def set_control_max_units(value):
+                                shared.opts.control_max_units = value
+                            control_max_units.change(fn=set_control_max_units, inputs=[control_max_units], outputs=[])
+                            control_tiles = gr.Textbox(label="Tiling options", value=shared.opts.control_tiles, elem_id='control_tiles')
+                            def set_control_tiles(value):
+                                shared.opts.control_tiles = value
+                            control_tiles.change(fn=set_control_tiles, inputs=[control_tiles], outputs=[])
+                            control_move_processor = gr.Checkbox(label="Move processor to CPU after use", value=shared.opts.control_move_processor, elem_id='control_move_processor')
+                            def set_control_move_processor(value):
+                                shared.opts.control_move_processor = value
+                            control_move_processor.change(fn=set_control_move_processor, inputs=[control_move_processor], outputs=[])
+                            control_unload_processor = gr.Checkbox(label="Unload processor after use", value=shared.opts.control_unload_processor, elem_id='control_unload_processor')
+                            def set_control_unload_processor(value):
+                                shared.opts.control_unload_processor = value
+                            control_unload_processor.change(fn=set_control_unload_processor, inputs=[control_unload_processor], outputs=[])
+
                         with gr.Accordion('HED', open=True, elem_classes=['processor-settings']):
                             settings.append(gr.Checkbox(label="Scribble", value=False))
                         with gr.Accordion('Midas depth', open=True, elem_classes=['processor-settings']):

modules/ui_extra_networks.py

@@ -941,9 +941,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         from modules.processing_info import get_last_args
         params, text = get_last_args()
         if (not params) or (not text) or (len(text) == 0):
-            filename = os.path.join(paths.data_path, "params.txt")
-            if os.path.exists(filename):
-                with open(filename, "r", encoding="utf8") as file:
+            if os.path.exists(paths.params_path):
+                with open(paths.params_path, "r", encoding="utf8") as file:
                     text = file.read()
             else:
                 text = ''
@@ -960,9 +959,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False):
         from modules.processing_info import get_last_args
         params, text = get_last_args()
         if (not params) or (not text) or (len(text) == 0):
-            fn = os.path.join(paths.data_path, "params.txt")
-            if os.path.exists(fn):
-                with open(fn, "r", encoding="utf8") as file:
+            if os.path.exists(paths.params_path):
+                with open(paths.params_path, "r", encoding="utf8") as file:
                     text = file.read()
             else:
                 text = ''

modules/video_models/video_load.py

@@ -1,4 +1,5 @@
 import os
+import copy
 import time
 from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices, sd_hijack_te
 from modules.video_models import models_def, video_utils, video_vae, video_overrides, video_cache
@@ -70,6 +71,9 @@ def load_model(selected: models_def.Model):
         errors.display(e, 'video')
 
     t1 = time.time()
+    if shared.sd_model.__class__.__name__.startswith("LTX"):
+        shared.sd_model.scheduler.config.use_dynamic_shifting = False
+    shared.sd_model.default_scheduler = copy.deepcopy(shared.sd_model.scheduler) if hasattr(shared.sd_model, "scheduler") else None
     shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo)
     shared.sd_model.sd_model_hash = None
     sd_models.set_diffuser_options(shared.sd_model)

repositories/codeformer/basicsr/__init__.py

@@ -8,4 +8,5 @@ from .models import *
 from .ops import *
 from .train import *
 from .utils import *
-from .version import __gitsha__, __version__
+__gitsha__ = '366a46c91d51923c56e09963dbc358bc61315408'
+__version__ = '1.3.2'

repositories/codeformer/basicsr/losses/losses.py

@@ -1,5 +1,4 @@
 import math
-import lpips
 import torch
 from torch import autograd as autograd
 from torch import nn as nn
@@ -260,6 +259,7 @@ class LPIPSLoss(nn.Module):
             use_input_norm=True,
             range_norm=False,):
         super(LPIPSLoss, self).__init__()
+        import lpips
         self.perceptual = lpips.LPIPS(net="vgg", spatial=False).eval()
         self.loss_weight = loss_weight
         self.use_input_norm = use_input_norm

requirements.txt

@@ -1,5 +1,6 @@
 # required for python 3.12
 setuptools==69.5.1
+wheel
 
 # standard
 patch-ng
@@ -33,7 +34,7 @@ pi-heif
 # versioned
 rich==14.0.0
 safetensors==0.5.3
-tensordict==0.1.2
+tensordict==0.8.3
 peft==0.15.2
 httpx==0.24.1
 compel==2.0.3
@@ -45,14 +46,15 @@ accelerate==1.7.0
 opencv-contrib-python-headless==4.9.0.80
 einops==0.4.1
 gradio==3.43.2
-huggingface_hub==0.31.2
+huggingface_hub==0.33.0
 numexpr==2.10.2
 numpy==1.26.4
+pandas==2.3.0
 numba==0.61.2
 protobuf==4.25.3
 pytorch_lightning==1.9.4
 tokenizers==0.21.1
-transformers==4.52.3
+transformers==4.52.4
 urllib3==1.26.19
 Pillow==10.4.0
 timm==0.9.16
@@ -63,7 +65,6 @@ typing-extensions==4.12.2
 # additional
 blendmodes
 scipy
-pandas
 torchdiffeq
 dctorch
 scikit-image

scripts/xyz_grid_classes.py

@@ -1,4 +1,4 @@
-from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
+from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, apply_sdnq_quant,  apply_sdnq_quant_te, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import
 from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet
 
 
@@ -58,6 +58,7 @@ class SharedSettingsStackHelper(object):
     extra_networks_default_multiplier = None
     disable_apply_metadata = None
     disable_apply_params = None
+    sdnq_quant_mode = None
 
     def __enter__(self):
         # Save overridden settings so they can be restored later
@@ -89,6 +90,8 @@ class SharedSettingsStackHelper(object):
         self.teacache_thresh = shared.opts.teacache_thresh
         self.disable_apply_metadata = shared.opts.disable_apply_metadata
         self.disable_apply_params = shared.opts.disable_apply_params
+        self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode
+
         shared.opts.data["disable_apply_metadata"] = []
         shared.opts.data["disable_apply_params"] = ''
 
@@ -135,6 +138,9 @@ class SharedSettingsStackHelper(object):
         if self.sd_unet != shared.opts.sd_unet:
             shared.opts.data["sd_unet"] = self.sd_unet
             sd_unet.load_unet(shared.sd_model)
+        if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode:
+            shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode
+            sd_models.reload_model_weights(op='model')
 
 
 axis_options = [
@@ -193,6 +199,8 @@ axis_options = [
     AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]),
     AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label),
     AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")),
+    AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
+    AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)),
     AxisOption("[HDR] Mode", int, apply_field("hdr_mode")),
     AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")),
     AxisOption("[HDR] Color", float, apply_field("hdr_color")),

scripts/xyz_grid_shared.py

@@ -147,6 +147,18 @@ def confirm_samplers(p, xs):
             shared.log.warning(f"XYZ grid: unknown sampler: {x}")
 
 
+def apply_sdnq_quant(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode = x
+    sd_models.unload_model_weights(op='model') # reload will happen on-demand
+    shared.log.debug(f'XYZ grid apply sdnq quant: mode="{x}"')
+
+
+def apply_sdnq_quant_te(p, x, xs):
+    shared.opts.sdnq_quantize_weights_mode_te = x
+    sd_models.unload_model_weights(op='model') # reload will happen on-demand
+    shared.log.debug(f'XYZ grid apply sdnq quant te: mode="{x}"')
+
+
 def apply_checkpoint(p, x, xs):
     if x == shared.opts.sd_model_checkpoint:
         return