nunchaku sdxl and sdxl-turbo support

Signed-off-by: Vladimir Mandic <mandic00@live.com>

CHANGELOG.md

@@ -7,6 +7,11 @@
     available for *text-to-image* and *text-to-video* and *image-to-video* workflows  
   - [Tencent FLUX.1 Dev SRPO](https://huggingface.co/tencent/SRPO)  
     SRPO is trained by with specific technique: Directly Aligning the Full Diffusion Trajectory with Fine-Grained Human Preference  
+  - [Nunchaku SDXL](https://huggingface.co/nunchaku-tech/nunchaku-sdxl) and [Nunchaku SDXL Turbo](https://huggingface.co/nunchaku-tech/nunchaku-sdxl-turbo)  
+    impact of nunchaku engine on unet-based model such as sdxl is much less than on a dit-based models, but its still significantly faster than baseline  
+    note that nunchaku optimized and prequantized unet is replacement for base unet, so its only applicable to base models, not any of finetunes  
+    *how to use*: enable nunchaku in settings -> quantization and then load either sdxl-base or sdxl-base-turbo reference models  
+    *note*: sdxl support for nunchaku is not in released version of `nunchaku==1.0.0`, so you need to build [nunchaku](https://nunchaku.tech/docs/nunchaku/installation/installation.html) from source
 - **Offloading**
   - improve offloading for pipelines with multiple stages such as *wan-2.2-14b*  
   - add timers to measure onload/offload times during generate  

html/reference.json

@@ -55,11 +55,21 @@
     "desc": "This stable-diffusion-2 model is resumed from stable-diffusion-2-base (512-base-ema.ckpt) and trained for 150k steps using a v-objective on the same dataset. Resumed for another 140k steps on 768x768 images",
     "extras": "width: 768, height: 768, sampler: DEIS, steps: 20, cfg_scale: 6.0"
   },
-  "StabilityAI StableDiffusion XL 1.0 Base": {
-    "path": "sd_xl_base_1.0.safetensors@https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/resolve/main/sd_xl_base_1.0.safetensors?download=true",
+  "StabilityAI StableDiffusion XL": {
+    "path": "stabilityai/stable-diffusion-xl-base-1.0",
     "preview": "sd_xl_base_1.0.jpg",
     "desc": "Stable Diffusion XL (SDXL) is the latest AI image generation model that is tailored towards more photorealistic outputs with more detailed imagery and composition compared to previous SD models, including SD 2.1. It can make realistic faces, legible text within the images, and better image composition, all while using shorter and simpler prompts at a greatly increased base resolution of 1024x1024. Just like its predecessors, SDXL has the ability to generate image variations using image-to-image prompting, inpainting (reimagining of the selected parts of an image), and outpainting (creating new parts that lie outside the image borders).",
-    "extras": "sampler: DEIS, steps: 20, cfg_scale: 6.0"
+    "skip": true,
+    "variant": "fp16",
+    "extras": ""
+  },
+  "StabilityAI StableDiffusion XL Turbo": {
+    "path": "stabilityai/sdxl-turbo",
+    "preview": "sd_xl_base_1.0.jpg",
+    "desc": "SDXL-Turbo is a fast generative text-to-image model that can synthesize photorealistic images from a text prompt in a single network evaluation.",
+    "skip": true,
+    "variant": "fp16",
+    "extras": ""
   },
   "StabilityAI Stable Cascade": {
     "path": "huggingface/stabilityai/stable-cascade",

modules/civitai/download_civitai.py

@@ -103,7 +103,7 @@ def download_civit_model_thread(model_name: str, model_url: str, model_path: str
     if os.path.isfile(temp_file):
         starting_pos = os.path.getsize(temp_file)
         headers['Range'] = f'bytes={starting_pos}-'
-    if ('civit' in model_url.lower()):
+    if 'civit' in model_url.lower(): # downloader can be used for other urls too
         if token is None or len(token) == 0:
             token = shared.opts.civitai_token
         if (token is not None) and (len(token) > 0):

modules/sd_detect.py

@@ -115,6 +115,8 @@ def guess_by_name(fn, current_guess):
         return 'Kandinsky 3.0'
     elif 'hunyuanimage' in fn.lower():
         return 'HunyuanImage'
+    elif 'sdxl-turbo' in fn.lower() or 'stable-diffusion-xl' in fn.lower():
+        return 'Stable Diffusion XL'
     return current_guess
 
 

modules/sd_models.py

@@ -411,65 +411,70 @@ def load_diffuser_folder(model_type, pipeline, checkpoint_info, diffusers_load_c
     files = shared.walk_files(checkpoint_info.path, ['.safetensors', '.bin', '.ckpt'])
     if 'variant' not in diffusers_load_config and any('diffusion_pytorch_model.fp16' in f for f in files): # deal with diffusers lack of variant fallback when loading
         diffusers_load_config['variant'] = 'fp16'
-    if (model_type is not None) and (pipeline is not None) and ('ONNX' in model_type): # forced pipeline
-        try:
-            sd_model = pipeline.from_pretrained(checkpoint_info.path)
-        except Exception as e:
-            shared.log.error(f'Load {op}: type=ONNX path="{checkpoint_info.path}" {e}')
-            if debug_load:
-                errors.display(e, 'Load')
-            return None
-    else:
-        err1, err2, err3 = None, None, None
-        if os.path.exists(checkpoint_info.path) and os.path.isdir(checkpoint_info.path):
-            if os.path.exists(os.path.join(checkpoint_info.path, 'unet', 'diffusion_pytorch_model.bin')):
-                shared.log.debug(f'Load {op}: type=pickle')
-                diffusers_load_config['use_safetensors'] = False
+
+    err0, err1, err2, err3 = None, None, None, None
+    if os.path.exists(checkpoint_info.path) and os.path.isdir(checkpoint_info.path):
+        if os.path.exists(os.path.join(checkpoint_info.path, 'unet', 'diffusion_pytorch_model.bin')):
+            shared.log.debug(f'Load {op}: type=pickle')
+            diffusers_load_config['use_safetensors'] = False
+    if debug_load:
+        shared.log.debug(f'Load {op}: args={diffusers_load_config}')
+
+    try: #0 - using detected model type and pipeline
+        if (model_type is not None) and (pipeline is not None):
+            sd_model = pipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            sd_model.model_type = sd_model.__class__.__name__
+    except Exception as e:
+        err0 = e
         if debug_load:
-            shared.log.debug(f'Load {op}: args={diffusers_load_config}')
-        try: # 1 - autopipeline, best choice but not all pipelines are available
-            try:
+            errors.display(e, 'Load Detected')
+
+    try: # 1 - autopipeline, best choice but not all pipelines are available
+        try:
+            if err0 is not None:
                 sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
                 sd_model.model_type = sd_model.__class__.__name__
-            except ValueError as e:
-                if 'no variant default' in str(e):
-                    shared.log.warning(f'Load {op}: variant={diffusers_load_config["variant"]} model="{checkpoint_info.path}" using default variant')
-                    diffusers_load_config.pop('variant', None)
-                    sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
-                    sd_model.model_type = sd_model.__class__.__name__
-                elif 'safetensors found in directory' in str(err1):
-                    shared.log.warning(f'Load {op}: type=pickle')
-                    diffusers_load_config['use_safetensors'] = False
-                    sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
-                    sd_model.model_type = sd_model.__class__.__name__
-                else:
-                    raise ValueError from e # reraise
-        except Exception as e:
-            err1 = e
-            if debug_load:
-                errors.display(e, 'Load AutoPipeline')
-            # shared.log.error(f'AutoPipeline: {e}')
-        try: # 2 - diffusion pipeline, works for most non-linked pipelines
-            if err1 is not None:
-                sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+        except ValueError as e:
+            if 'no variant default' in str(e):
+                shared.log.warning(f'Load {op}: variant={diffusers_load_config["variant"]} model="{checkpoint_info.path}" using default variant')
+                diffusers_load_config.pop('variant', None)
+                sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
                 sd_model.model_type = sd_model.__class__.__name__
-        except Exception as e:
-            err2 = e
-            if debug_load:
-                errors.display(e, "Load DiffusionPipeline")
-            # shared.log.error(f'DiffusionPipeline: {e}')
-        try: # 3 - try basic pipeline just in case
-            if err2 is not None:
-                sd_model = diffusers.StableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            elif 'safetensors found in directory' in str(err1):
+                shared.log.warning(f'Load {op}: type=pickle')
+                diffusers_load_config['use_safetensors'] = False
+                sd_model = diffusers.AutoPipelineForText2Image.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
                 sd_model.model_type = sd_model.__class__.__name__
-        except Exception as e:
-            err3 = e  # ignore last error
-            shared.log.error(f"StableDiffusionPipeline: {e}")
-            if debug_load:
-                errors.display(e, "Load StableDiffusionPipeline")
-        if err3 is not None:
-            shared.log.error(f'Load {op}: {checkpoint_info.path} auto={err1} diffusion={err2}')
-            return None
+            else:
+                raise ValueError from e # reraise
+    except Exception as e:
+        err1 = e
+        if debug_load:
+            errors.display(e, 'Load AutoPipeline')
+
+    try: # 2 - diffusion pipeline, works for most non-linked pipelines
+        if err1 is not None:
+            sd_model = diffusers.DiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            sd_model.model_type = sd_model.__class__.__name__
+    except Exception as e:
+        err2 = e
+        if debug_load:
+            errors.display(e, "Load DiffusionPipeline")
+
+    try: # 3 - try basic pipeline just in case
+        if err2 is not None:
+            sd_model = diffusers.StableDiffusionPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+            sd_model.model_type = sd_model.__class__.__name__
+    except Exception as e:
+        err3 = e  # ignore last error
+        shared.log.error(f"StableDiffusionPipeline: {e}")
+        if debug_load:
+            errors.display(e, "Load StableDiffusionPipeline")
+
+    if err3 is not None:
+        shared.log.error(f'Load {op}: {checkpoint_info.path} detected={err0} auto={err1} diffusion={err2} base={err3}')
+        return None
+
     return sd_model
 
 
@@ -667,7 +672,7 @@ def load_diffuser(checkpoint_info=None, op='model', revision=None): # pylint: di
             sd_model.scheduler.name = 'DDIM'
 
         if hasattr(sd_model, "unet") and model_type not in ['Stable Cascade']: # others calls load_diffuser again
-            sd_unet.load_unet(sd_model)
+            sd_unet.load_unet(sd_model, checkpoint_info.path)
 
         add_noise_pred_to_diffusers_callback(sd_model)
 
@@ -1029,7 +1034,13 @@ def set_diffusers_attention(pipe, quiet:bool=False):
             return
         # other models uses their own attention processor
         if pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet"):
-            pipe.unet.set_attn_processor(attention)
+            try:
+                pipe.unet.set_attn_processor(attention)
+            except Exception as e:
+                if 'Nunchaku' in pipe.unet.__class__.__name__:
+                    pass
+                else:
+                    shared.log.error(f"Attention: {name if name is not None else attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}")
         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__}")
 

modules/sd_offload.py

@@ -293,7 +293,7 @@ def get_module_names(pipe=None, exclude=[]):
         modules_names = get_signature(pipe).keys()
     modules_names = [m for m in modules_names if m not in exclude and not m.startswith('_')]
     modules_names = [m for m in modules_names if isinstance(getattr(pipe, m, None), torch.nn.Module)]
-    modules_names = list(sorted(set(modules_names)))
+    modules_names = sorted(set(modules_names))
     return modules_names
 
 

modules/sd_unet.py

@@ -1,5 +1,5 @@
 import os
-from modules import shared, devices, files_cache, sd_models
+from modules import shared, devices, files_cache, sd_models, model_quant
 
 
 unet_dict = {}
@@ -8,22 +8,55 @@ failed_unet = []
 debug = os.environ.get('SD_LOAD_DEBUG', None) is not None
 
 
-dit_models = ['Flux', 'StableDiffusion3', 'HiDream', 'Lumina2', 'Chroma', 'Wan']
+dit_models = ['Flux', 'StableDiffusion3', 'HiDream', 'Lumina2', 'Chroma', 'Wan', 'Qwen']
 
 
-def load_unet(model):
+def load_unet_sdxl_nunchaku(repo_id):
+    try:
+        from nunchaku.models.unets.unet_sdxl import NunchakuSDXLUNet2DConditionModel
+    except Exception:
+        shared.log.error(f'Load module: quant=Nunchaku module=unet repo="{repo_id}" low nunchaku version')
+        return None
+    if 'turbo' in repo_id.lower():
+        nunchaku_repo = 'nunchaku-tech/nunchaku-sdxl-turbo/svdq-int4_r32-sdxl-turbo.safetensors'
+    else:
+        nunchaku_repo = 'nunchaku-tech/nunchaku-sdxl/svdq-int4_r32-sdxl.safetensors'
+
+    shared.log.debug(f'Load module: quant=Nunchaku module=unet repo="{nunchaku_repo}" offload={shared.opts.nunchaku_offload}')
+    unet = NunchakuSDXLUNet2DConditionModel.from_pretrained(
+        nunchaku_repo,
+        offload=shared.opts.nunchaku_offload,
+        torch_dtype=devices.dtype,
+        cache_dir=shared.opts.hfcache_dir,
+    )
+    unet.quantization_method = 'SVDQuant'
+    return unet
+
+
+def load_unet(model, repo_id:str=None):
     global loaded_unet # pylint: disable=global-statement
+
+    if ("StableDiffusionXLPipeline" in model.__class__.__name__) and (('stable-diffusion-xl-base' in repo_id) or ('sdxl-turbo' in repo_id)):
+        if model_quant.check_nunchaku('Model'):
+            unet = load_unet_sdxl_nunchaku(repo_id)
+            if unet is not None:
+                model.unet = unet
+                return
+
     if shared.opts.sd_unet == 'Default' or shared.opts.sd_unet == 'None':
         return
+
     if shared.opts.sd_unet not in list(unet_dict):
         shared.log.error(f'Load module: type=UNet not found: {shared.opts.sd_unet}')
         return
+
     config_file = os.path.splitext(unet_dict[shared.opts.sd_unet])[0] + '.json'
     if os.path.exists(config_file):
         config = shared.readfile(config_file)
     else:
         config = None
         config_file = 'default'
+
     try:
         if shared.opts.sd_unet == loaded_unet or shared.opts.sd_unet in failed_unet:
             pass

modules/sdnq/common.py

@@ -1,8 +1,8 @@
 # pylint: disable=redefined-builtin,no-member,protected-access
 
 import os
-import torch
 from functools import partial
+import torch
 
 from modules import shared
 
@@ -46,5 +46,5 @@ if use_torch_compile:
     torch._dynamo.config.accumulated_recompile_limit = max(8192, torch._dynamo.config.accumulated_recompile_limit)
     compile_func = partial(torch.compile, fullgraph=True, dynamic=False)
 else:
-    def compile_func(fn, **kwargs):
+    def compile_func(fn, **kwargs): # pylint: disable=unused-argument
         return fn

pipelines/flux/flux_legacy_loader.py

@@ -130,7 +130,7 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unus
                 shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported')
             if nunchaku_repo is not None:
                 shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}')
-                kwargs['transformer'] = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype)
+                kwargs['transformer'] = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype, cache_dir=cache_dir)
                 kwargs['transformer'].quantization_method = 'SVDQuant'
                 if shared.opts.nunchaku_attention:
                     kwargs['transformer'].set_attention_impl("nunchaku-fp16")
@@ -142,7 +142,7 @@ def load_quants(kwargs, repo_id, cache_dir, allow_quant): # pylint: disable=unus
             nunchaku_precision = nunchaku.utils.get_precision()
             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)
+            kwargs['text_encoder_2'] = nunchaku.NunchakuT5EncoderModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype, cache_dir=cache_dir)
             kwargs['text_encoder_2'].quantization_method = 'SVDQuant'
         if 'text_encoder_2' not in kwargs and model_quant.check_quant('TE'):
             load_args, quant_args = model_quant.get_dit_args(diffusers_load_config, module='TE', device_map=True)

pipelines/flux/flux_nunchaku.py

@@ -26,7 +26,12 @@ def load_flux_nunchaku(repo_id):
         shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported')
     if nunchaku_repo is not None:
         shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}')
-        transformer = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype)
+        transformer = nunchaku.NunchakuFluxTransformer2dModel.from_pretrained(
+            nunchaku_repo,
+            offload=shared.opts.nunchaku_offload,
+            torch_dtype=devices.dtype,
+            cache_dir=shared.opts.hfcache_dir,
+        )
         transformer.quantization_method = 'SVDQuant'
         if shared.opts.nunchaku_attention:
             transformer.set_attention_impl("nunchaku-fp16")

pipelines/model_qwen.py

@@ -26,8 +26,6 @@ def load_qwen(checkpoint_info, diffusers_load_config={}):
     if model_quant.check_nunchaku('Model'):
         from pipelines.qwen.qwen_nunchaku import load_qwen_nunchaku
         transformer = load_qwen_nunchaku(repo_id)
-        # if transformer is not None:
-        #     cls_name = nunchaku.pipeline.pipeline_qwenimage.NunchakuQwenImagePipeline # we dont need this
 
     if 'Qwen-Image-Distill-Full' in repo_id:
         repo_transformer = repo_id

pipelines/model_sana.py

@@ -10,9 +10,8 @@ def load_quants(kwargs, repo_id, cache_dir):
         import nunchaku
         nunchaku_precision = nunchaku.utils.get_precision()
         nunchaku_repo = "nunchaku-tech/nunchaku-sana/svdq-int4_r32-sana1.6b.safetensors"
-        # https://huggingface.co/nunchaku-tech/nunchaku-sana/blob/main/svdq-int4_r32-sana1.6b.safetensors
         shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} attention={shared.opts.nunchaku_attention}')
-        kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype)
+        kwargs['transformer'] = nunchaku.NunchakuSanaTransformer2DModel.from_pretrained(nunchaku_repo, torch_dtype=devices.dtype, cache_dir=cache_dir)
     elif model_quant.check_quant('Model'):
         load_args, quant_args = model_quant.get_dit_args(kwargs_copy, module='Model')
         kwargs['transformer'] = diffusers.SanaTransformer2DModel.from_pretrained(repo_id, subfolder="transformer", cache_dir=cache_dir, **load_args, **quant_args)

pipelines/qwen/qwen_nunchaku.py

@@ -23,6 +23,11 @@ def load_qwen_nunchaku(repo_id):
         shared.log.error(f'Load module: quant=Nunchaku module=transformer repo="{repo_id}" unsupported')
     if nunchaku_repo is not None:
         shared.log.debug(f'Load module: quant=Nunchaku module=transformer repo="{nunchaku_repo}" precision={nunchaku_precision} offload={shared.opts.nunchaku_offload} attention={shared.opts.nunchaku_attention}')
-        transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(nunchaku_repo, offload=shared.opts.nunchaku_offload, torch_dtype=devices.dtype) # pylint: disable=no-member
+        transformer = NunchakuQwenImageTransformer2DModel.from_pretrained(
+            nunchaku_repo,
+            offload=shared.opts.nunchaku_offload,
+            torch_dtype=devices.dtype,
+            cache_dir=shared.opts.hfcache_dir,
+        ) # pylint: disable=no-member
         transformer.quantization_method = 'SVDQuant'
     return transformer