nunchaku sdxl and sdxl-turbo support

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

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