Merge branch 'dev' into Dev/DirectoryCacheing

modules/cmd_args.py

@@ -80,12 +80,12 @@ def compatibility_args(opts, args):
     group.add_argument("--swinir-models-path", help=argparse.SUPPRESS, default=opts.swinir_models_path)
     group.add_argument("--ldsr-models-path", help=argparse.SUPPRESS, default=opts.ldsr_models_path)
     group.add_argument("--clip-models-path", type=str, help=argparse.SUPPRESS, default=opts.clip_models_path)
-    group.add_argument("--opt-channelslast", help=argparse.SUPPRESS, default=opts.opt_channelslast)
+    group.add_argument("--opt-channelslast", help=argparse.SUPPRESS, action='store_true', default=opts.opt_channelslast)
     group.add_argument("--xformers", default=(opts.cross_attention_optimization == "xFormers"), action='store_true', help=argparse.SUPPRESS)
-    group.add_argument("--disable-nan-check", help=argparse.SUPPRESS, default=opts.disable_nan_check)
+    group.add_argument("--disable-nan-check", help=argparse.SUPPRESS, action='store_true', default=opts.disable_nan_check)
     group.add_argument("--rollback-vae", help=argparse.SUPPRESS, default=opts.rollback_vae)
-    group.add_argument("--no-half", help=argparse.SUPPRESS, default=opts.no_half)
-    group.add_argument("--no-half-vae", help=argparse.SUPPRESS, default=opts.no_half_vae)
+    group.add_argument("--no-half", help=argparse.SUPPRESS, action='store_true', default=opts.no_half)
+    group.add_argument("--no-half-vae", help=argparse.SUPPRESS, action='store_true', default=opts.no_half_vae)
     group.add_argument("--precision", help=argparse.SUPPRESS, default=opts.precision)
     group.add_argument("--sub-quad-q-chunk-size", help=argparse.SUPPRESS, default=opts.sub_quad_q_chunk_size)
     group.add_argument("--sub-quad-kv-chunk-size", help=argparse.SUPPRESS, default=opts.sub_quad_kv_chunk_size)
@@ -94,7 +94,7 @@ def compatibility_args(opts, args):
     group.add_argument("--lyco-dir", help=argparse.SUPPRESS, default=opts.lyco_dir)
     group.add_argument("--embeddings-dir", help=argparse.SUPPRESS, default=opts.embeddings_dir)
     group.add_argument("--hypernetwork-dir", help=argparse.SUPPRESS, default=opts.hypernetwork_dir)
-    group.add_argument("--lyco-patch-lora", help=argparse.SUPPRESS, default=False)
+    group.add_argument("--lyco-patch-lora", help=argparse.SUPPRESS, action='store_true', default=False)
     group.add_argument("--lyco-debug", help=argparse.SUPPRESS, action='store_true', default=False)
     group.add_argument("--enable-console-prompts", help=argparse.SUPPRESS, action='store_true', default=False)
     group.add_argument("--safe", help=argparse.SUPPRESS, action='store_true', default=False)

modules/control/run.py

@@ -467,11 +467,9 @@ def control_run(units: List[unit.Unit], inputs, inits, mask, unit_type: str, is_
                     # pipeline
                     output = None
                     if pipe is not None: # run new pipeline
-                        debug(f'Control exec pipeline: class={pipe.__class__}')
-                        debug(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)}')
+                        debug(f'Control exec pipeline: task={sd_models.get_diffusers_task(pipe)} class={pipe.__class__}')
                         debug(f'Control exec pipeline: p={vars(p)}')
-                        debug(f'Control exec pipeline: args={p.task_args}')
-                        debug(f'Control exec pipeline: image={p.task_args.get("image", None)} control={p.task_args.get("control_image", None)} mask={p.task_args.get("mask_image", None)} ref={p.task_args.get("ref_image", None)}')
+                        debug(f'Control exec pipeline: args={p.task_args} image={p.task_args.get("image", None)} control={p.task_args.get("control_image", None)} mask={p.task_args.get("mask_image", None)} ref={p.task_args.get("ref_image", None)}')
                         processed: processing.Processed = processing.process_images(p) # run actual pipeline
                         output = processed.images if processed is not None else None
                         # output = pipe(**vars(p)).images # alternative direct pipe exec call

modules/control/unit.py

@@ -1,5 +1,6 @@
 from typing import Union
 from PIL import Image
+import gradio as gr
 from modules.shared import log
 from modules.control import processors
 from modules.control.units import controlnet
@@ -31,6 +32,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                  image_input = None,
                  preview_process = None,
                  image_upload = None,
+                 image_preview = None,
                  control_start = None,
                  control_end = None,
                  result_txt = None,
@@ -101,8 +103,10 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
                 self.process.override = Image.open(image_file.name)
                 self.override = self.process.override
                 log.debug(f'Control process upload image: path="{image_file.name}" image={self.process.override}')
+                return gr.update(visible=self.process.override is not None, value=self.process.override)
             except Exception as e:
                 log.error(f'Control process upload image failed: path="{image_file.name}" error={e}')
+                return gr.update(visible=False, value=None)
 
         # actual init
         if self.type == 'adapter':
@@ -157,7 +161,7 @@ class Unit(): # mashup of gradio controls and mapping to actual implementation c
         if preview_btn is not None:
             preview_btn.click(fn=self.process.preview, inputs=[self.input], outputs=[preview_process]) # return list of images for gallery
         if image_upload is not None:
-            image_upload.upload(fn=upload_image, inputs=[image_upload], outputs=[]) # return list of images for gallery
+            image_upload.upload(fn=upload_image, inputs=[image_upload], outputs=[image_preview]) # return list of images for gallery
         if control_start is not None and control_end is not None:
             control_start.change(fn=control_change, inputs=[control_start, control_end])
             control_end.change(fn=control_change, inputs=[control_start, control_end])

modules/control/units/controlnet.py

@@ -3,7 +3,7 @@ import time
 from typing import Union
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, ControlNetModel, StableDiffusionControlNetPipeline, StableDiffusionXLControlNetPipeline
 from modules.control.units import detect
-from modules.shared import log, opts
+from modules.shared import log, opts, listdir
 from modules import errors
 
 
@@ -44,7 +44,7 @@ cache_dir = 'models/control/controlnet'
 
 def find_models():
     path = os.path.join(opts.control_dir, 'controlnet')
-    files = os.listdir(path)
+    files = listdir(path)
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:

modules/control/units/lite.py

@@ -4,7 +4,7 @@ from typing import Union
 import numpy as np
 from PIL import Image
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
-from modules.shared import log, opts
+from modules.shared import log, opts, listdir
 from modules import errors
 from modules.control.units.lite_model import ControlNetLLLite
 
@@ -31,7 +31,7 @@ cache_dir = 'models/control/lite'
 
 def find_models():
     path = os.path.join(opts.control_dir, 'lite')
-    files = os.listdir(path)
+    files = listdir(path)
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:

modules/control/units/reference.py

@@ -3,7 +3,7 @@ import time
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from modules.control.proc.reference_sd15 import StableDiffusionReferencePipeline
 from modules.control.proc.reference_sdxl import StableDiffusionXLReferencePipeline
-from modules.shared import log
+from modules.shared import log, opts
 from modules.control.units import detect
 
 
@@ -22,6 +22,8 @@ class ReferencePipeline():
         if pipeline is None:
             log.error(f'Control {what} model pipeline: model not loaded')
             return
+        if opts.diffusers_fuse_projections and hasattr(pipeline, 'unfuse_qkv_projections'):
+            pipeline.unfuse_qkv_projections()
         if detect.is_sdxl(pipeline):
             self.pipeline = StableDiffusionXLReferencePipeline(
                 vae=pipeline.vae,

modules/control/units/xs.py

@@ -2,7 +2,7 @@ import os
 import time
 from typing import Union
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
-from modules.shared import log, opts
+from modules.shared import log, opts, listdir
 from modules import errors
 from modules.control.units.xs_model import ControlNetXSModel
 from modules.control.units.xs_pipe import StableDiffusionControlNetXSPipeline, StableDiffusionXLControlNetXSPipeline
@@ -27,7 +27,7 @@ cache_dir = 'models/control/xs'
 
 def find_models():
     path = os.path.join(opts.control_dir, 'xs')
-    files = os.listdir(path)
+    files = listdir(path)
     files = [f for f in files if f.endswith('.safetensors')]
     downloaded_models = {}
     for f in files:

modules/images.py

@@ -504,7 +504,7 @@ def get_next_sequence_number(path, basename):
     prefix_length = len(basename)
     if not os.path.isdir(path):
         return 0
-    for p in os.listdir(path):
+    for p in shared.listdir(path):
         if p.startswith(basename):
             parts = os.path.splitext(p[prefix_length:])[0].split('-')  # splits the filename (removing the basename first if one is defined, so the sequence number is always the first element)
             try:

modules/img2img.py

@@ -22,10 +22,10 @@ def process_batch(p, input_files, input_dir, output_dir, inpaint_mask_dir, args)
         if not os.path.isdir(input_dir):
             shared.log.error(f"Process batch: directory not found: {input_dir}")
             return
-        image_files = shared.listfiles(input_dir)
+        image_files = shared.listdir(input_dir)
     is_inpaint_batch = False
     if inpaint_mask_dir:
-        inpaint_masks = shared.listfiles(inpaint_mask_dir)
+        inpaint_masks = shared.listdir(inpaint_mask_dir)
         is_inpaint_batch = len(inpaint_masks) > 0
     if is_inpaint_batch:
         shared.log.info(f"Process batch: inpaint batch masks={len(inpaint_masks)}")

modules/intel/ipex/hijacks.py

@@ -1,4 +1,5 @@
-import contextlib
+from functools import wraps
+from contextlib import nullcontext
 import torch
 import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
 from modules import devices
@@ -12,7 +13,7 @@ class DummyDataParallel(torch.nn.Module): # pylint: disable=missing-class-docstr
         return module.to(devices.device)
 
 def return_null_context(*args, **kwargs): # pylint: disable=unused-argument
-    return contextlib.nullcontext()
+    return nullcontext()
 
 @property
 def is_cuda(self):
@@ -27,6 +28,7 @@ def return_xpu(device):
 
 # Autocast
 original_autocast = torch.autocast
+@wraps(torch.autocast)
 def ipex_autocast(*args, **kwargs):
     if len(args) > 0 and (args[0] == "cuda" or args[0] == "xpu"):
         if "dtype" in kwargs:
@@ -38,6 +40,7 @@ def ipex_autocast(*args, **kwargs):
 
 # Latent Antialias CPU Offload:
 original_interpolate = torch.nn.functional.interpolate
+@wraps(torch.nn.functional.interpolate)
 def interpolate(tensor, size=None, scale_factor=None, mode='nearest', align_corners=None, recompute_scale_factor=None, antialias=False): # pylint: disable=too-many-arguments
     if antialias or align_corners is not None:
         return_device = tensor.device
@@ -50,6 +53,7 @@ def interpolate(tensor, size=None, scale_factor=None, mode='nearest', align_corn
 
 # Diffusers Float64 (Alchemist GPUs doesn't support 64 bit):
 original_from_numpy = torch.from_numpy
+@wraps(torch.from_numpy)
 def from_numpy(ndarray):
     if ndarray.dtype == float:
         return original_from_numpy(ndarray.astype('float32'))
@@ -70,11 +74,13 @@ else:
 
 
 # Data Type Errors:
+@wraps(torch.bmm)
 def torch_bmm(input, mat2, *, out=None):
     if input.dtype != mat2.dtype:
         mat2 = mat2.to(input.dtype)
     return original_torch_bmm(input, mat2, out=out)
 
+@wraps(torch.nn.functional.scaled_dot_product_attention)
 def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False):
     if query.dtype != key.dtype:
         key = key.to(dtype=query.dtype)
@@ -84,6 +90,7 @@ def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.
 
 # A1111 FP16
 original_functional_group_norm = torch.nn.functional.group_norm
+@wraps(torch.nn.functional.group_norm)
 def functional_group_norm(input, num_groups, weight=None, bias=None, eps=1e-05):
     if weight is not None and input.dtype != weight.data.dtype:
         input = input.to(dtype=weight.data.dtype)
@@ -93,6 +100,7 @@ def functional_group_norm(input, num_groups, weight=None, bias=None, eps=1e-05):
 
 # A1111 BF16
 original_functional_layer_norm = torch.nn.functional.layer_norm
+@wraps(torch.nn.functional.layer_norm)
 def functional_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
     if weight is not None and input.dtype != weight.data.dtype:
         input = input.to(dtype=weight.data.dtype)
@@ -102,6 +110,7 @@ def functional_layer_norm(input, normalized_shape, weight=None, bias=None, eps=1
 
 # Training
 original_functional_linear = torch.nn.functional.linear
+@wraps(torch.nn.functional.linear)
 def functional_linear(input, weight, bias=None):
     if input.dtype != weight.data.dtype:
         input = input.to(dtype=weight.data.dtype)
@@ -110,6 +119,7 @@ def functional_linear(input, weight, bias=None):
     return original_functional_linear(input, weight, bias=bias)
 
 original_functional_conv2d = torch.nn.functional.conv2d
+@wraps(torch.nn.functional.conv2d)
 def functional_conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
     if input.dtype != weight.data.dtype:
         input = input.to(dtype=weight.data.dtype)
@@ -119,6 +129,7 @@ def functional_conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1,
 
 # A1111 Embedding BF16
 original_torch_cat = torch.cat
+@wraps(torch.cat)
 def torch_cat(tensor, *args, **kwargs):
     if len(tensor) == 3 and (tensor[0].dtype != tensor[1].dtype or tensor[2].dtype != tensor[1].dtype):
         return original_torch_cat([tensor[0].to(tensor[1].dtype), tensor[1], tensor[2].to(tensor[1].dtype)], *args, **kwargs)
@@ -127,6 +138,7 @@ def torch_cat(tensor, *args, **kwargs):
 
 # SwinIR BF16:
 original_functional_pad = torch.nn.functional.pad
+@wraps(torch.nn.functional.pad)
 def functional_pad(input, pad, mode='constant', value=None):
     if mode == 'reflect' and input.dtype == torch.bfloat16:
         return original_functional_pad(input.to(torch.float32), pad, mode=mode, value=value).to(dtype=torch.bfloat16)
@@ -135,6 +147,7 @@ def functional_pad(input, pad, mode='constant', value=None):
 
 
 original_torch_tensor = torch.tensor
+@wraps(torch.tensor)
 def torch_tensor(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_tensor(*args, device=return_xpu(device), **kwargs)
@@ -142,6 +155,7 @@ def torch_tensor(*args, device=None, **kwargs):
         return original_torch_tensor(*args, device=device, **kwargs)
 
 original_Tensor_to = torch.Tensor.to
+@wraps(torch.Tensor.to)
 def Tensor_to(self, device=None, *args, **kwargs):
     if check_device(device):
         return original_Tensor_to(self, return_xpu(device), *args, **kwargs)
@@ -149,6 +163,7 @@ def Tensor_to(self, device=None, *args, **kwargs):
         return original_Tensor_to(self, device, *args, **kwargs)
 
 original_Tensor_cuda = torch.Tensor.cuda
+@wraps(torch.Tensor.cuda)
 def Tensor_cuda(self, device=None, *args, **kwargs):
     if check_device(device):
         return original_Tensor_cuda(self, return_xpu(device), *args, **kwargs)
@@ -156,6 +171,7 @@ def Tensor_cuda(self, device=None, *args, **kwargs):
         return original_Tensor_cuda(self, device, *args, **kwargs)
 
 original_UntypedStorage_init = torch.UntypedStorage.__init__
+@wraps(torch.UntypedStorage.__init__)
 def UntypedStorage_init(*args, device=None, **kwargs):
     if check_device(device):
         return original_UntypedStorage_init(*args, device=return_xpu(device), **kwargs)
@@ -163,6 +179,7 @@ def UntypedStorage_init(*args, device=None, **kwargs):
         return original_UntypedStorage_init(*args, device=device, **kwargs)
 
 original_UntypedStorage_cuda = torch.UntypedStorage.cuda
+@wraps(torch.UntypedStorage.cuda)
 def UntypedStorage_cuda(self, device=None, *args, **kwargs):
     if check_device(device):
         return original_UntypedStorage_cuda(self, return_xpu(device), *args, **kwargs)
@@ -170,6 +187,7 @@ def UntypedStorage_cuda(self, device=None, *args, **kwargs):
         return original_UntypedStorage_cuda(self, device, *args, **kwargs)
 
 original_torch_empty = torch.empty
+@wraps(torch.empty)
 def torch_empty(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_empty(*args, device=return_xpu(device), **kwargs)
@@ -177,6 +195,7 @@ def torch_empty(*args, device=None, **kwargs):
         return original_torch_empty(*args, device=device, **kwargs)
 
 original_torch_randn = torch.randn
+@wraps(torch.randn)
 def torch_randn(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_randn(*args, device=return_xpu(device), **kwargs)
@@ -184,6 +203,7 @@ def torch_randn(*args, device=None, **kwargs):
         return original_torch_randn(*args, device=device, **kwargs)
 
 original_torch_ones = torch.ones
+@wraps(torch.ones)
 def torch_ones(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_ones(*args, device=return_xpu(device), **kwargs)
@@ -191,6 +211,7 @@ def torch_ones(*args, device=None, **kwargs):
         return original_torch_ones(*args, device=device, **kwargs)
 
 original_torch_zeros = torch.zeros
+@wraps(torch.zeros)
 def torch_zeros(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_zeros(*args, device=return_xpu(device), **kwargs)
@@ -198,6 +219,7 @@ def torch_zeros(*args, device=None, **kwargs):
         return original_torch_zeros(*args, device=device, **kwargs)
 
 original_torch_linspace = torch.linspace
+@wraps(torch.linspace)
 def torch_linspace(*args, device=None, **kwargs):
     if check_device(device):
         return original_torch_linspace(*args, device=return_xpu(device), **kwargs)
@@ -205,6 +227,7 @@ def torch_linspace(*args, device=None, **kwargs):
         return original_torch_linspace(*args, device=device, **kwargs)
 
 original_torch_Generator = torch.Generator
+@wraps(torch.Generator)
 def torch_Generator(device=None):
     if check_device(device):
         return original_torch_Generator(return_xpu(device))
@@ -212,6 +235,7 @@ def torch_Generator(device=None):
         return original_torch_Generator(device)
 
 original_torch_load = torch.load
+@wraps(torch.load)
 def torch_load(f, map_location=None, pickle_module=None, *, weights_only=False, mmap=None, **kwargs):
     if check_device(map_location):
         return original_torch_load(f, map_location=return_xpu(map_location), pickle_module=pickle_module, weights_only=weights_only, mmap=mmap, **kwargs)

modules/intel/openvino/__init__.py

@@ -229,8 +229,8 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na
         om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
-    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights):
-        if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8":
+    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compile_dont_use_4bit and not shared.opts.nncf_compress_vae_weights):
+        if shared.compiled_model_state.compile_dont_use_4bit or shared.opts.nncf_compress_weights_mode == "INT8":
             om = nncf.compress_weights(om)
         else:
             om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito)
@@ -238,7 +238,7 @@ def openvino_compile(gm: GraphModule, *args, model_hash_str: str = None, file_na
     if model_hash_str is not None:
         core.set_property({'CACHE_DIR': cache_root + '/blob'})
 
-    shared.compiled_model_state.compiling_vae = False
+    shared.compiled_model_state.compile_dont_use_4bit = False
     compiled_model = core.compile_model(om, device)
     return compiled_model
 
@@ -261,15 +261,15 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
         om.inputs[idx].get_node().set_element_type(dtype_mapping[input_data.dtype])
         om.inputs[idx].get_node().set_partial_shape(PartialShape(list(input_data.shape)))
     om.validate_nodes_and_infer_types()
-    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compiling_vae and not shared.opts.nncf_compress_vae_weights):
-        if shared.compiled_model_state.compiling_vae or shared.opts.nncf_compress_weights_mode == "INT8":
+    if shared.opts.nncf_compress_weights and not (shared.compiled_model_state.compile_dont_use_4bit and not shared.opts.nncf_compress_vae_weights):
+        if shared.compiled_model_state.compile_dont_use_4bit or shared.opts.nncf_compress_weights_mode == "INT8":
             om = nncf.compress_weights(om)
         else:
             om = nncf.compress_weights(om, mode=getattr(nncf.CompressWeightsMode, shared.opts.nncf_compress_weights_mode), group_size=8, ratio=shared.opts.nncf_compress_weights_raito)
 
     core.set_property({'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'})
 
-    shared.compiled_model_state.compiling_vae = False
+    shared.compiled_model_state.compile_dont_use_4bit = False
     compiled_model = core.compile_model(om, get_device())
     return compiled_model
 
@@ -344,7 +344,11 @@ def partition_graph(gm: GraphModule, use_python_fusion_cache: bool, model_hash_s
 
 def generate_subgraph_str(tensor):
     if hasattr(tensor, "weight"):
-        shared.compiled_model_state.model_str = shared.compiled_model_state.model_str + str(tensor.weight)
+        shared.compiled_model_state.model_str = shared.compiled_model_state.model_str + sha256(str(tensor.weight).encode('utf-8')).hexdigest()
+    return tensor
+
+def get_subgraph_type(tensor):
+    shared.compiled_model_state.subgraph_type.append(type(tensor))
     return tensor
 
 @register_backend
@@ -353,12 +357,28 @@ def openvino_fx(subgraph, example_inputs):
     executor_parameters = None
     inputs_reversed = False
     maybe_fs_cached_name = None
+
+    shared.compiled_model_state.subgraph_type = []
+    subgraph.apply(get_subgraph_type)
+
+    # SD 1.5 / SDXL VAE
+    if (shared.compiled_model_state.subgraph_type[0] is torch.nn.modules.conv.Conv2d and
+        shared.compiled_model_state.subgraph_type[1] is torch.nn.modules.conv.Conv2d and
+        shared.compiled_model_state.subgraph_type[2] is torch.nn.modules.normalization.GroupNorm and
+        shared.compiled_model_state.subgraph_type[3] is torch.nn.modules.activation.SiLU):
+
+        shared.compiled_model_state.compile_dont_use_4bit = True
+
     if not shared.opts.openvino_disable_model_caching:
         os.environ.setdefault('OPENVINO_TORCH_MODEL_CACHING', "1")
-        # Create a hash to be used for caching
         shared.compiled_model_state.model_str = ""
+
+        # Create a hash to be used for caching
         subgraph.apply(generate_subgraph_str)
+        shared.compiled_model_state.model_str = shared.compiled_model_state.model_str + sha256(subgraph.code.encode('utf-8')).hexdigest()
         model_hash_str = sha256(shared.compiled_model_state.model_str.encode('utf-8')).hexdigest()
+        shared.compiled_model_state.model_str = ""
+
         if (shared.compiled_model_state.cn_model != [] and shared.compiled_model_state.partition_id == 0):
             model_hash_str = model_hash_str + str(shared.compiled_model_state.cn_model)
 
@@ -383,9 +403,17 @@ def openvino_fx(subgraph, example_inputs):
                                 example_inputs_reordered.append(example_inputs[idx1])
                 example_inputs = example_inputs_reordered
 
-            # Delete unused subgraphs
-            subgraph = subgraph.apply(sd_models.convert_to_faketensors)
-            devices.torch_gc(force=True)
+            # SD 1.5 / SDXL Text Encoder
+            if (shared.compiled_model_state.subgraph_type[0] is torch.nn.modules.sparse.Embedding and
+                shared.compiled_model_state.subgraph_type[1] is torch.nn.modules.sparse.Embedding and
+                shared.compiled_model_state.subgraph_type[2] is torch.nn.modules.normalization.LayerNorm and
+                shared.compiled_model_state.subgraph_type[3] is torch.nn.modules.linear.Linear):
+
+                pass # Fails with FakeTensors or Downcast
+            else:
+                # Delete unused subgraphs
+                subgraph = subgraph.apply(sd_models.convert_to_faketensors)
+                devices.torch_gc(force=True)
 
             # Model is fully supported and already cached. Run the cached OV model directly.
             compiled_model = openvino_compile_cached_model(maybe_fs_cached_name, *example_inputs)

modules/modelloader.py

@@ -230,7 +230,7 @@ def load_diffusers_models(model_path: str, command_path: str = None, clear=True)
                 if not os.path.isfile(os.path.join(cache_path, "hidden")):
                     output.append(str(r.repo_id))
             """
-            for folder in os.listdir(place):
+            for folder in shared.listdir(place):
                 try:
                     if "--" not in folder:
                         continue
@@ -240,7 +240,7 @@ def load_diffusers_models(model_path: str, command_path: str = None, clear=True)
                     name = name.replace("--", "/")
                     folder = os.path.join(place, folder)
                     friendly = os.path.join(place, name)
-                    snapshots = os.listdir(os.path.join(folder, "snapshots"))
+                    snapshots = shared.listdir(os.path.join(folder, "snapshots"))
                     if len(snapshots) == 0:
                         shared.log.warning(f"Diffusers folder has no snapshots: location={place} folder={folder} name={name}")
                         continue
@@ -450,7 +450,7 @@ def move_files(src_path: str, dest_path: str, ext_filter: str = None):
         if not os.path.exists(dest_path):
             os.makedirs(dest_path)
         if os.path.exists(src_path):
-            for file in os.listdir(src_path):
+            for file in shared.listdir(src_path):
                 fullpath = os.path.join(src_path, file)
                 if os.path.isfile(fullpath):
                     if ext_filter is not None:

modules/postprocessing.py

@@ -38,7 +38,7 @@ def run_postprocessing(extras_mode, image, image_folder: List[tempfile.NamedTemp
     elif extras_mode == 2:
         assert not shared.cmd_opts.hide_ui_dir_config, '--hide-ui-dir-config option must be disabled'
         assert input_dir, 'input directory not selected'
-        image_list = shared.listfiles(input_dir)
+        image_list = shared.listdir(input_dir)
         for filename in image_list:
             try:
                 image = Image.open(filename)

modules/processing_vae.py

@@ -48,13 +48,9 @@ def full_vae_decode(latents, model):
         model.upcast_vae()
         latents = latents.to(next(iter(model.vae.post_quant_conv.parameters())).dtype)
 
-    # OpenVINO with INT4 doesn't work with VAE decode so we pass that we are using VAE right now to OpenVINO
-    if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae:
-        shared.compiled_model_state.compiling_vae = True
-
     decoded = model.vae.decode(latents / model.vae.config.scaling_factor, return_dict=False)[0]
 
-    # Downcast VAE after OpenVINO compile
+    # Delete PyTorch VAE after OpenVINO compile
     if shared.opts.cuda_compile and shared.opts.cuda_compile_backend == "openvino_fx" and shared.compiled_model_state.first_pass_vae:
         shared.compiled_model_state.first_pass_vae = False
         if hasattr(shared.sd_model, "vae"):

modules/script_loading.py

@@ -7,6 +7,7 @@ from installer import setup_logging, args
 
 
 preloaded = []
+debug = os.environ.get('SD_SCRIPT_DEBUG', None)
 
 
 def load_module(path):
@@ -20,9 +21,12 @@ def load_module(path):
         if '/sd-extension-' in path: # safe extensions without stdout intercept
             module_spec.loader.exec_module(module)
         else:
-            # stdout = io.StringIO()
-            with contextlib.redirect_stdout(io.StringIO()) as stdout:
+            if debug:
                 module_spec.loader.exec_module(module)
+                stdout = io.StringIO()
+            else:
+                with contextlib.redirect_stdout(io.StringIO()) as stdout:
+                    module_spec.loader.exec_module(module)
             setup_logging() # reset since scripts can hijaack logging
             for line in stdout.getvalue().splitlines():
                 if len(line) > 0:

modules/sd_hijack_hypertile.py

@@ -190,7 +190,7 @@ def context_hypertile_vae(p):
         # shared.log.warning('Hypertile VAE is enabled but no VAE model was found')
         return nullcontext()
     else:
-        tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128))
+        tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(128, 64 * min(p.width // 128, p.height // 128))
         shared.log.info(f'Applying hypertile: vae={tile_size}')
         p.extra_generation_params['Hypertile VAE'] = tile_size
         return split_attention(vae, tile_size=tile_size, min_tile_size=128, swap_size=1)
@@ -216,7 +216,7 @@ def context_hypertile_unet(p):
         # shared.log.warning('Hypertile UNet is enabled but no Unet model was found')
         return nullcontext()
     else:
-        tile_size = shared.opts.hypertile_unet_tile if shared.opts.hypertile_unet_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128))
+        tile_size = shared.opts.hypertile_unet_tile if shared.opts.hypertile_unet_tile > 0 else max(128, 64 * min(p.width // 128, p.height // 128))
         shared.log.info(f'Applying hypertile: unet={tile_size}')
         p.extra_generation_params['Hypertile UNet'] = tile_size
         return split_attention(unet, tile_size=tile_size, min_tile_size=128, swap_size=1)

modules/sd_models.py

@@ -1036,6 +1036,10 @@ def set_diffuser_pipe(pipe, new_pipe_type):
     image_encoder = getattr(pipe, "image_encoder", None)
     feature_extractor = getattr(pipe, "feature_extractor", None)
 
+    # skip specific pipelines
+    if pipe.__class__.__name__ == 'StableDiffusionReferencePipeline' or pipe.__class__.__name__ == 'StableDiffusionAdapterPipeline':
+        return pipe
+
     try:
         if new_pipe_type == DiffusersTaskType.TEXT_2_IMAGE:
             new_pipe = diffusers.AutoPipelineForText2Image.from_pipe(pipe)
@@ -1044,7 +1048,7 @@ def set_diffuser_pipe(pipe, new_pipe_type):
         elif new_pipe_type == DiffusersTaskType.INPAINTING:
             new_pipe = diffusers.AutoPipelineForInpainting.from_pipe(pipe)
     except Exception as e: # pylint: disable=unused-variable
-        shared.log.error(f'Failed to change: type={new_pipe_type} pipeline={pipe.__class__.__name__} {e}')
+        shared.log.warning(f'Failed to change: type={new_pipe_type} pipeline={pipe.__class__.__name__} {e}')
         return pipe
 
     if pipe.__class__ == new_pipe.__class__:
@@ -1263,10 +1267,10 @@ def reload_model_weights(sd_model=None, info=None, reuse_dict=False, op='model')
 
 
 def convert_to_faketensors(tensor):
-        fake = torch._subclasses.fake_tensor.FakeTensorMode()
-        if hasattr(tensor, "weight"):
-            tensor.weight = torch.nn.Parameter(fake.from_tensor(tensor.weight))
-        return tensor
+    fake_module = torch._subclasses.fake_tensor.FakeTensorMode(allow_non_fake_inputs=True) # pylint: disable=protected-access
+    if hasattr(tensor, "weight"):
+        tensor.weight = torch.nn.Parameter(fake_module.from_tensor(tensor.weight))
+    return tensor
 
 
 def disable_offload(sd_model):
@@ -1280,6 +1284,9 @@ def disable_offload(sd_model):
 
 
 def unload_model_weights(op='model', change_from='none'):
+    if shared.compiled_model_state is not None:
+        shared.compiled_model_state.compiled_cache.clear()
+        shared.compiled_model_state.partitioned_modules.clear()
     if op == 'model' or op == 'dict':
         if model_data.sd_model:
             if (shared.backend == shared.Backend.ORIGINAL and change_from != shared.Backend.DIFFUSERS) or change_from == shared.Backend.ORIGINAL:

modules/sd_models_compile.py

@@ -21,7 +21,8 @@ class CompiledModelState:
         self.lora_compile = False
         self.compiled_cache = {}
         self.partitioned_modules = {}
-        self.compiling_vae = False
+        self.subgraph_type = []
+        self.compile_dont_use_4bit = False
 
 
 def ipex_optimize(sd_model):
@@ -72,13 +73,13 @@ def optimize_openvino():
         if shared.compiled_model_state is None:
             shared.compiled_model_state = CompiledModelState()
         else:
-            if not shared.compiled_model_state.lora_compile:
-                shared.compiled_model_state.lora_compile = False
-                shared.compiled_model_state.lora_model = []
             shared.compiled_model_state.compiled_cache.clear()
             shared.compiled_model_state.partitioned_modules.clear()
-            shared.compiled_model_state.partition_id = 0
-            shared.compiled_model_state.model_str = ""
+            backup_lora_model = []
+            if shared.compiled_model_state.lora_compile:
+                backup_lora_model = shared.compiled_model_state.lora_model
+            shared.compiled_model_state = CompiledModelState()
+            shared.compiled_model_state.lora_model = backup_lora_model
         shared.compiled_model_state.first_pass = True if not shared.opts.cuda_compile_precompile else False
         shared.compiled_model_state.first_pass_vae = True if not shared.opts.cuda_compile_precompile else False
         shared.compiled_model_state.first_pass_refiner = True if not shared.opts.cuda_compile_precompile else False
@@ -149,10 +150,10 @@ def compile_torch(sd_model):
 
         t0 = time.time()
         if shared.opts.cuda_compile:
-            if shared.opts.cuda_compile and (not hasattr(sd_model, 'unet') or not hasattr(sd_model.unet, 'config')):
-                shared.log.warning('Model compile enabled but model has no Unet')
-            else:
+            if hasattr(sd_model, 'unet') and hasattr(sd_model.unet, 'config'):
                 sd_model.unet = torch.compile(sd_model.unet, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+            else:
+                shared.log.warning('Model compile enabled but model has no Unet')
         if shared.opts.cuda_compile_vae:
             if hasattr(sd_model, 'vae') and hasattr(sd_model.vae, 'decode'):
                 sd_model.vae.decode = torch.compile(sd_model.vae.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
@@ -160,6 +161,13 @@ def compile_torch(sd_model):
                 sd_model.movq.decode = torch.compile(sd_model.movq.decode, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
             else:
                 shared.log.warning('Model compile enabled but model has no VAE')
+        if shared.opts.cuda_compile_text_encoder:
+            if hasattr(sd_model, 'text_encoder'):
+                sd_model.text_encoder = torch.compile(sd_model.text_encoder, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+                if hasattr(sd_model, 'text_encoder_2'):
+                    sd_model.text_encoder_2 = torch.compile(sd_model.text_encoder_2, mode=shared.opts.cuda_compile_mode, backend=shared.opts.cuda_compile_backend, fullgraph=shared.opts.cuda_compile_fullgraph)
+            else:
+                shared.log.warning('Text Encoder compile enabled but model has no Text Encoder')
         setup_logging() # compile messes with logging so reset is needed
         if shared.opts.cuda_compile_precompile:
             sd_model("dummy prompt")

modules/shared.py

@@ -70,6 +70,8 @@ restricted_opts = {
 resize_modes = ["None", "Fixed", "Crop", "Fill", "Latent"]
 compatibility_opts = ['clip_skip', 'uni_pc_lower_order_final', 'uni_pc_order']
 console = Console(log_time=True, log_time_format='%H:%M:%S-%f')
+dir_timestamps = {}
+dir_cache = {}
 
 
 class Backend(Enum):
@@ -327,7 +329,8 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
     "cuda_compile_sep": OptionInfo("<h2>Model Compile</h2>", "", gr.HTML),
     "cuda_compile": OptionInfo(False if not cmd_opts.use_openvino else True, "Compile UNet"),
     "cuda_compile_vae": OptionInfo(False if not cmd_opts.use_openvino else True, "Compile VAE"),
-    "cuda_compile_upscaler": OptionInfo(False if not cmd_opts.use_openvino else True, "Compile upscaler"),
+    "cuda_compile_text_encoder": OptionInfo(False, "Compile Text Encoder"),
+    "cuda_compile_upscaler": OptionInfo(False if not cmd_opts.use_openvino else True, "Compile Upscaler"),
     "cuda_compile_backend": OptionInfo("none" if not cmd_opts.use_openvino else "openvino_fx", "Model compile backend", gr.Radio, {"choices": ['none', 'inductor', 'cudagraphs', 'aot_ts_nvfuser', 'hidet', 'ipex', 'openvino_fx', 'stable-fast']}),
     "cuda_compile_mode": OptionInfo("default", "Model compile mode", gr.Radio, {"choices": ['default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs']}),
     "cuda_compile_fullgraph": OptionInfo(False, "Model compile fullgraph"),
@@ -887,9 +890,16 @@ def restore_defaults(restart=True):
     restart_server(restart)
 
 
-def listfiles(dirname):
-    filenames = [os.path.join(dirname, x) for x in sorted(os.listdir(dirname), key=str.lower) if not x.startswith(".")]
-    return [file for file in filenames if os.path.isfile(file)]
+def listdir(path):
+    if not os.path.exists(path):
+        return []
+    mtime = os.path.getmtime(path)
+    if path in dir_timestamps and mtime == dir_timestamps[path]:
+        return dir_cache[path]
+    else:
+        dir_cache[path] = [os.path.join(path, f) for f in os.listdir(path)]
+        dir_timestamps[path] = mtime
+        return dir_cache[path]
 
 
 def walk_files(path, allowed_extensions=None):

modules/textual_inversion/preprocess.py

@@ -25,10 +25,6 @@ def preprocess(id_task, process_src, process_dst, process_width, process_height,
             deepbooru.model.stop()
 
 
-def listfiles(dirname):
-    return os.listdir(dirname)
-
-
 class PreprocessParams:
     src = None
     dstdir = None
@@ -137,17 +133,12 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
     dst = os.path.abspath(process_dst)
     split_threshold = max(0.0, min(1.0, split_threshold))
     overlap_ratio = max(0.0, min(0.9, overlap_ratio))
-
     assert src != dst, 'same directory specified as source and destination'
-
     os.makedirs(dst, exist_ok=True)
-
-    files = listfiles(src)
-
+    files = shared.listdir(src)
     shared.state.job = "preprocess"
     shared.state.textinfo = "Preprocessing..."
     shared.state.job_count = len(files)
-
     params = PreprocessParams()
     params.dstdir = dst
     params.flip = process_flip
@@ -155,7 +146,6 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
     params.process_caption = process_caption
     params.process_caption_deepbooru = process_caption_deepbooru
     params.preprocess_txt_action = preprocess_txt_action
-
     pbar = tqdm(files)
     for index, imagefile in enumerate(pbar):
         params.subindex = 0
@@ -171,9 +161,7 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
         description = f"Preprocessing image {index + 1}/{len(files)}"
         pbar.set_description(description)
         shared.state.textinfo = description
-
         params.src = filename
-
         existing_caption = None
         existing_caption_filename = f"{os.path.splitext(filename)[0]}.txt"
         if os.path.exists(existing_caption_filename):
@@ -181,32 +169,25 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
                 existing_caption = file.read()
         else:
             existing_caption_filename = None
-
         if shared.state.interrupted:
             break
-
         if img.height > img.width:
             ratio = (img.width * height) / (img.height * width)
             inverse_xy = False
         else:
             ratio = (img.height * width) / (img.width * height)
             inverse_xy = True
-
         process_default_resize = True
-
         if process_split and ratio < 1.0 and ratio <= split_threshold:
             for splitted in split_pic(img, inverse_xy, width, height, overlap_ratio):
                 save_pic(splitted, index, params, existing_caption=existing_caption, existing_caption_filename=existing_caption_filename)
             process_default_resize = False
-
         if process_focal_crop and img.height != img.width:
-
             dnn_model_path = None
             try:
                 dnn_model_path = autocrop.download_and_cache_models(os.path.join(paths.models_path, "opencv"))
             except Exception as e:
                 print("Unable to load face detection model for auto crop selection. Falling back to lower quality haar method.", e)
-
             autocrop_settings = autocrop.Settings(
                 crop_width = width,
                 crop_height = height,
@@ -227,13 +208,10 @@ def preprocess_work(process_src, process_dst, process_width, process_height, pre
             else:
                 print(f"skipped {img.width}x{img.height} image {filename} (can't find suitable size within error threshold)")
             process_default_resize = False
-
         if process_keep_original_size:
             save_pic(img, index, params, existing_caption=existing_caption)
             process_default_resize = False
-
         if process_default_resize:
             img = images.resize_image(1, img, width, height)
             save_pic(img, index, params, existing_caption=existing_caption)
-
         shared.state.nextjob()

modules/ui_control.py

@@ -421,20 +421,19 @@ def create_ui(_blocks: gr.Blocks=None):
                         num_controlnet_units = gr.Slider(label="Units", minimum=1, maximum=max_units, step=1, value=1, scale=1)
                     controlnet_ui_units = [] # list of hidable accordions
                     for i in range(max_units):
-                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_controlnet_units.value) as unit_ui:
+                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_controlnet_units.value, elem_classes='control-unit') as unit_ui:
                             with gr.Row():
-                                with gr.Column():
-                                    with gr.Row():
-                                        enabled_cb = gr.Checkbox(value= i==0, label="")
-                                        process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                        model_id = gr.Dropdown(label="ControlNet", choices=controlnet.list_models(), value='None')
-                                        ui_common.create_refresh_button(model_id, controlnet.list_models, lambda: {"choices": controlnet.list_models(refresh=True)}, 'refresh_controlnet_models')
-                                        model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
-                                        control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
-                                        control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
-                                        reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
-                                        image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
-                                        process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                enabled_cb = gr.Checkbox(value= i==0, label="")
+                                process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
+                                model_id = gr.Dropdown(label="ControlNet", choices=controlnet.list_models(), value='None')
+                                ui_common.create_refresh_button(model_id, controlnet.list_models, lambda: {"choices": controlnet.list_models(refresh=True)}, 'refresh_controlnet_models')
+                                model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
+                                control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
+                                control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
+                                reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
+                                image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
+                                process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
                         controlnet_ui_units.append(unit_ui)
                         units.append(unit.Unit(
                             unit_type = 'controlnet',
@@ -448,6 +447,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             preview_process = preview_process,
                             preview_btn = process_btn,
                             image_upload = image_upload,
+                            image_preview = image_preview,
                             control_start = control_start,
                             control_end = control_end,
                             extra_controls = extra_controls,
@@ -457,51 +457,6 @@ def create_ui(_blocks: gr.Blocks=None):
                             units[-1].enabled = True # enable first unit in group
                     num_controlnet_units.change(fn=display_units, inputs=[num_controlnet_units], outputs=controlnet_ui_units)
 
-                with gr.Tab('XS') as _tab_controlnetxs:
-                    gr.HTML('<a href="https://vislearn.github.io/ControlNet-XS/">ControlNet XS</a>')
-                    with gr.Row():
-                        extra_controls = [
-                            gr.Slider(label="Time embedding mix", minimum=0.0, maximum=1.0, step=0.05, value=0.0, scale=3)
-                        ]
-                        num_controlnet_units = gr.Slider(label="Units", minimum=1, maximum=max_units, step=1, value=1, scale=1)
-                    controlnetxs_ui_units = [] # list of hidable accordions
-                    for i in range(max_units):
-                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_controlnet_units.value) as unit_ui:
-                            with gr.Row():
-                                with gr.Column():
-                                    with gr.Row():
-                                        enabled_cb = gr.Checkbox(value= i==0, label="")
-                                        process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                        model_id = gr.Dropdown(label="ControlNet-XS", choices=xs.list_models(), value='None')
-                                        ui_common.create_refresh_button(model_id, xs.list_models, lambda: {"choices": xs.list_models(refresh=True)}, 'refresh_xs_models')
-                                        model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
-                                        control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
-                                        control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
-                                        reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
-                                        image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
-                                        process_btn= ui_components.ToolButton(value=ui_symbols.preview)
-                        controlnetxs_ui_units.append(unit_ui)
-                        units.append(unit.Unit(
-                            unit_type = 'xs',
-                            result_txt = result_txt,
-                            image_input = input_image,
-                            enabled_cb = enabled_cb,
-                            reset_btn = reset_btn,
-                            process_id = process_id,
-                            model_id = model_id,
-                            model_strength = model_strength,
-                            preview_process = preview_process,
-                            preview_btn = process_btn,
-                            image_upload = image_upload,
-                            control_start = control_start,
-                            control_end = control_end,
-                            extra_controls = extra_controls,
-                            )
-                        )
-                        if i == 0:
-                            units[-1].enabled = True # enable first unit in group
-                    num_controlnet_units.change(fn=display_units, inputs=[num_controlnet_units], outputs=controlnetxs_ui_units)
-
                 with gr.Tab('Adapter') as _tab_adapter:
                     gr.HTML('<a href="https://github.com/TencentARC/T2I-Adapter">T2I-Adapter</a>')
                     with gr.Row():
@@ -511,18 +466,17 @@ def create_ui(_blocks: gr.Blocks=None):
                         num_adapter_units = gr.Slider(label="Units", minimum=1, maximum=max_units, step=1, value=1, scale=1)
                     adapter_ui_units = [] # list of hidable accordions
                     for i in range(max_units):
-                        with gr.Accordion(f'Adapter unit {i+1}', visible= i < num_adapter_units.value) as unit_ui:
+                        with gr.Accordion(f'Adapter unit {i+1}', visible= i < num_adapter_units.value, elem_classes='control-unit') as unit_ui:
                             with gr.Row():
-                                with gr.Column():
-                                    with gr.Row():
-                                        enabled_cb = gr.Checkbox(value= i == 0, label="Enabled")
-                                        process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                        model_id = gr.Dropdown(label="Adapter", choices=t2iadapter.list_models(), value='None')
-                                        ui_common.create_refresh_button(model_id, t2iadapter.list_models, lambda: {"choices": t2iadapter.list_models(refresh=True)}, 'refresh_adapter_models')
-                                        model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
-                                        reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
-                                        image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
-                                        process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                enabled_cb = gr.Checkbox(value= i == 0, label="Enabled")
+                                process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
+                                model_id = gr.Dropdown(label="Adapter", choices=t2iadapter.list_models(), value='None')
+                                ui_common.create_refresh_button(model_id, t2iadapter.list_models, lambda: {"choices": t2iadapter.list_models(refresh=True)}, 'refresh_adapter_models')
+                                model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
+                                reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
+                                image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
+                                process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
                         adapter_ui_units.append(unit_ui)
                         units.append(unit.Unit(
                             unit_type = 'adapter',
@@ -536,6 +490,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             preview_process = preview_process,
                             preview_btn = process_btn,
                             image_upload = image_upload,
+                            image_preview = image_preview,
                             extra_controls = extra_controls,
                             )
                         )
@@ -543,6 +498,51 @@ def create_ui(_blocks: gr.Blocks=None):
                             units[-1].enabled = True # enable first unit in group
                     num_adapter_units.change(fn=display_units, inputs=[num_adapter_units], outputs=adapter_ui_units)
 
+                with gr.Tab('XS') as _tab_controlnetxs:
+                    gr.HTML('<a href="https://vislearn.github.io/ControlNet-XS/">ControlNet XS</a>')
+                    with gr.Row():
+                        extra_controls = [
+                            gr.Slider(label="Time embedding mix", minimum=0.0, maximum=1.0, step=0.05, value=0.0, scale=3)
+                        ]
+                        num_controlnet_units = gr.Slider(label="Units", minimum=1, maximum=max_units, step=1, value=1, scale=1)
+                    controlnetxs_ui_units = [] # list of hidable accordions
+                    for i in range(max_units):
+                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_controlnet_units.value, elem_classes='control-unit') as unit_ui:
+                            with gr.Row():
+                                enabled_cb = gr.Checkbox(value= i==0, label="")
+                                process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
+                                model_id = gr.Dropdown(label="ControlNet-XS", choices=xs.list_models(), value='None')
+                                ui_common.create_refresh_button(model_id, xs.list_models, lambda: {"choices": xs.list_models(refresh=True)}, 'refresh_xs_models')
+                                model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
+                                control_start = gr.Slider(label="Start", minimum=0.0, maximum=1.0, step=0.05, value=0)
+                                control_end = gr.Slider(label="End", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
+                                reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
+                                image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
+                                process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                        controlnetxs_ui_units.append(unit_ui)
+                        units.append(unit.Unit(
+                            unit_type = 'xs',
+                            result_txt = result_txt,
+                            image_input = input_image,
+                            enabled_cb = enabled_cb,
+                            reset_btn = reset_btn,
+                            process_id = process_id,
+                            model_id = model_id,
+                            model_strength = model_strength,
+                            preview_process = preview_process,
+                            preview_btn = process_btn,
+                            image_upload = image_upload,
+                            image_preview = image_preview,
+                            control_start = control_start,
+                            control_end = control_end,
+                            extra_controls = extra_controls,
+                            )
+                        )
+                        if i == 0:
+                            units[-1].enabled = True # enable first unit in group
+                    num_controlnet_units.change(fn=display_units, inputs=[num_controlnet_units], outputs=controlnetxs_ui_units)
+
                 with gr.Tab('Lite') as _tab_lite:
                     gr.HTML('<a href="https://huggingface.co/kohya-ss/controlnet-lllite">Control LLLite</a>')
                     with gr.Row():
@@ -551,18 +551,17 @@ def create_ui(_blocks: gr.Blocks=None):
                         num_lite_units = gr.Slider(label="Units", minimum=1, maximum=max_units, step=1, value=1, scale=1)
                     lite_ui_units = [] # list of hidable accordions
                     for i in range(max_units):
-                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_lite_units.value) as unit_ui:
+                        with gr.Accordion(f'Control unit {i+1}', visible= i < num_lite_units.value, elem_classes='control-unit') as unit_ui:
                             with gr.Row():
-                                with gr.Column():
-                                    with gr.Row():
-                                        enabled_cb = gr.Checkbox(value= i == 0, label="Enabled")
-                                        process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
-                                        model_id = gr.Dropdown(label="Model", choices=lite.list_models(), value='None')
-                                        ui_common.create_refresh_button(model_id, lite.list_models, lambda: {"choices": lite.list_models(refresh=True)}, 'refresh_lite_models')
-                                        model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
-                                        reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
-                                        image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
-                                        process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                enabled_cb = gr.Checkbox(value= i == 0, label="Enabled")
+                                process_id = gr.Dropdown(label="Processor", choices=processors.list_models(), value='None')
+                                model_id = gr.Dropdown(label="Model", choices=lite.list_models(), value='None')
+                                ui_common.create_refresh_button(model_id, lite.list_models, lambda: {"choices": lite.list_models(refresh=True)}, 'refresh_lite_models')
+                                model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0-i/10)
+                                reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
+                                image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
+                                image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                process_btn= ui_components.ToolButton(value=ui_symbols.preview)
                         lite_ui_units.append(unit_ui)
                         units.append(unit.Unit(
                             unit_type = 'lite',
@@ -576,6 +575,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             preview_process = preview_process,
                             preview_btn = process_btn,
                             image_upload = image_upload,
+                            image_preview = image_preview,
                             extra_controls = extra_controls,
                             )
                         )
@@ -593,16 +593,15 @@ def create_ui(_blocks: gr.Blocks=None):
                             gr.Slider(label="Reference adain weight", minimum=0.0, maximum=2.0, step=0.05, value=1.0, interactive=True),
                         ]
                     for i in range(1): # can only have one reference unit
-                        with gr.Accordion(f'Reference unit {i+1}', visible=True) as unit_ui:
+                        with gr.Accordion(f'Reference unit {i+1}', visible=True, elem_classes='control-unit') as unit_ui:
                             with gr.Row():
-                                with gr.Column():
-                                    with gr.Row():
-                                        enabled_cb = gr.Checkbox(value= i == 0, label="Enabled", visible=False)
-                                        model_id = gr.Dropdown(label="Reference", choices=reference.list_models(), value='Reference', visible=False)
-                                        model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, visible=False)
-                                        reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
-                                        image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
-                                        process_btn= ui_components.ToolButton(value=ui_symbols.preview)
+                                enabled_cb = gr.Checkbox(value= i == 0, label="Enabled", visible=False)
+                                model_id = gr.Dropdown(label="Reference", choices=reference.list_models(), value='Reference', visible=False)
+                                model_strength = gr.Slider(label="Strength", minimum=0.01, maximum=1.0, step=0.01, value=1.0, visible=False)
+                                reset_btn = ui_components.ToolButton(value=ui_symbols.reset)
+                                image_upload = gr.UploadButton(label=ui_symbols.upload, file_types=['image'], elem_classes=['form', 'gradio-button', 'tool'])
+                                image_preview = gr.Image(label="Input", show_label=False, type="pil", source="upload", interactive=False, height=128, width=128, visible=False)
+                                process_btn= ui_components.ToolButton(value=ui_symbols.preview)
                         units.append(unit.Unit(
                             unit_type = 'reference',
                             result_txt = result_txt,
@@ -615,6 +614,7 @@ def create_ui(_blocks: gr.Blocks=None):
                             preview_process = preview_process,
                             preview_btn = process_btn,
                             image_upload = image_upload,
+                            image_preview = image_preview,
                             extra_controls = extra_controls,
                             )
                         )

modules/ui_img2img.py

@@ -20,7 +20,7 @@ def process_interrogate(interrogation_function, mode, ii_input_files, ii_input_d
             if not os.path.isdir(ii_input_dir):
                 shared.log.error(f"Interrogate: Input directory not found: {ii_input_dir}")
                 return [gr.update(), None]
-            images = shared.listfiles(ii_input_dir)
+            images = shared.listdir(ii_input_dir)
         if ii_output_dir != "":
             os.makedirs(ii_output_dir, exist_ok=True)
         else:

modules/ui_interrogate.py

@@ -119,7 +119,7 @@ def batch_process(batch_files, batch_folder, batch_str, mode, clip_model, write)
     if batch_folder is not None:
         files += [f.name for f in batch_folder]
     if batch_str is not None and len(batch_str) > 0 and os.path.exists(batch_str) and os.path.isdir(batch_str):
-        files += [os.path.join(batch_str, f) for f in os.listdir(batch_str) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
+        files += [os.path.join(batch_str, f) for f in shared.listdir(batch_str) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.webp'))]
     if len(files) == 0:
         shared.log.error('Interrogate batch no images')
         return ''

modules/upscaler.py

@@ -52,7 +52,7 @@ class Upscaler:
             pass
 
     def find_folder(self, folder, scalers, loaded):
-        for fn in os.listdir(folder): # from folder
+        for fn in modules.shared.listdir(folder): # from folder
             file_name = os.path.join(folder, fn)
             if os.path.isdir(file_name):
                 self.find_folder(file_name, scalers, loaded)
@@ -83,19 +83,6 @@ class Upscaler:
         if not os.path.exists(self.user_path):
             return scalers
         self.find_folder(self.user_path, scalers, loaded)
-        """
-        for fn in os.listdir(self.user_path): # from folder
-            if not fn.endswith('.pth') and not fn.endswith('.pt'):
-                continue
-            file_name = os.path.join(self.user_path, fn)
-            if file_name not in loaded:
-                model_name = os.path.splitext(fn)[0]
-                scaler = UpscalerData(name=f'{self.name} {model_name}', path=file_name, upscaler=self)
-                scaler.custom = True
-                scalers.append(scaler)
-                loaded.append(file_name)
-                # modules.shared.log.debug(f'Upscaler type={self.name} folder="{self.user_path}" model="{model_name}" path="{file_name}"')
-        """
         return scalers
 
     @abstractmethod