Merge branch 'dev' into bundled-emb

.github/ISSUE_TEMPLATE/bug_report.yml

@@ -26,14 +26,15 @@ body:
         Easiest is to include top part of console log, for example:  
         ```log
         Starting SD.Next
-        Python 3.10.6 on Linux
-        Version: abd7d160 Sat Jun 10 07:37:42 2023 -0400
-        nVidia CUDA toolkit detected
-        Torch 2.1.0.dev20230519+cu121
-        Torch backend: nVidia CUDA 12.1 cuDNN 8801
-        Torch detected GPU: NVIDIA GeForce RTX 3060 VRAM 12288 Arch (8, 6) Cores 28
-        Enabled extensions-builtin: [...]
-        Enabled extensions: [...]
+        Version: app=sd.next updated=2024-06-28 hash=1fc20e72 branch=dev url=https://github.com/vladmandic/automatic/tree/dev ui=dev
+        Branch sync failed: sdnext=dev ui=dev
+        Platform: arch=x86_64 cpu=x86_64 system=Linux release=5.15.153.1-microsoft-standard-WSL2 python=3.12.3
+        Torch allocator: "garbage_collection_threshold:0.80,max_split_size_mb:512"
+        Load packages: {'torch': '2.3.1+cu121', 'diffusers': '0.29.1', 'gradio': '3.43.2'}
+        Engine: backend=Backend.DIFFUSERS compute=cuda device=cuda attention="Scaled-Dot-Product" mode=no_grad
+        Device: device=NVIDIA GeForce RTX 4090 n=1 arch=sm_90 cap=(8, 9) cuda=12.1 cudnn=8902 driver=555.99
+        Extensions: enabled=['sd-webui-agent-scheduler', 'sd-extension-chainner', 'sd-extension-system-info', 'sdnext-modernui', 'Lora'] extensions-builtin
+        Extensions: enabled=[] extensions
         ```
   - type: markdown
     attributes:
@@ -73,6 +74,18 @@ body:
       default: 0
     validations:
       required: true
+  - type: dropdown
+    id: ui
+    attributes:
+      label: UI
+      description: Which UI are you're using?
+      options:
+        - None
+        - Standard
+        - ModernUI
+      default: 1
+    validations:
+      required: true
   - type: dropdown
     id: branch
     attributes:
@@ -90,11 +103,12 @@ body:
       label: Model
       description: What is the model type you're using?
       options:
-        - SD 1.5
-        - SD 2.1
-        - SD-XL
+        - StableDiffusion 1.5
+        - StableDiffusion 2.1
+        - StableDiffusion XL
+        - StableDiffusion 3
         - PixArt
-        - Stable Cascade
+        - StableCascade
         - Kandinsky
         - Other
       default: 0

CHANGELOG.md

@@ -1,21 +1,33 @@
 # Change Log for SD.Next
 
-## Update for 2024-06-21
+## Update for 2024-06-28
 
-### Highlights for 2024-06-21
+- enable `florence` VLM for all platforms, thanks @lshqqytiger!  
+- fix executing extensions with zero params  
+- fix nncf for lora, thanks @Disty0!
+- fix diffusers version detection for SD3
+- fix current step for higher order samplers
+- add SD3 with FP16 T5 to list of detected models
+- multiple ModernUI fixes
 
-Following zero-day **SD3** release, a week later here's a refresh with 10+ improvements  
+## Update for 2024-06-23
+
+### Highlights for 2024-06-23
+
+Following zero-day **SD3** release, a 10 days later here's a refresh with 10+ improvements  
 including full prompt attention, support for compressed weights, additional text-encoder quantization modes.  
 
 But there's more than SD3:  
-- support for quantized **T5** text encoder in all models that use T5: FP4/FP8/FP16/INT8 (SD3, PixArt-Σ, etc)  
+- support for quantized **T5** text encoder *FP16/FP8/FP4/INT8* in all models that use T5: SD3, PixArt-Σ, etc.  
 - support for **PixArt-Sigma** in small/medium/large variants  
 - support for **HunyuanDiT 1.1**  
+- additional **NNCF weights compression** support: SD3, PixArt, ControlNet, Lora  
+- integration of **MS Florence** VLM/VQA *Base* and *Large* models  
 - (finally) new release of **Torch-DirectML**  
-- additional efficiencies for users with low vram gpus  
+- additional efficiencies for users with low VRAM GPUs  
 - over 20 overall fixes  
 
-### Model Improvements
+### Model Improvements for 2024-06-23
 
 - **SD3**: enable tiny-VAE (TAESD) preview and non-full quality mode  
 - SD3: enable base LoRA support  
@@ -38,14 +50,18 @@ But there's more than SD3:
   *note* by default pixart-Σ uses full fp16 t5 encoder with large memory footprint  
   simply select in *settings -> model -> text encoder* before or after model load  
 - **HunyuanDiT**: support for model version 1.1  
+- **MS Florence**: integration of Microsoft Florence VLM/VQA Base and Large models  
+  simply select in *process -> visual query*!
 
-### Improvements: General
+### General Improvements for 2024-06-23
 
 - support FP4 quantized T5 text encoder, in addtion to existing FP8 and FP16
 - support for T5 text-encoder loader in **all** models that use T5  
   *example*: load FP4 or FP8 quantized T5 text-encoder into PixArt Sigma!
 - support for `torch-directml` **0.2.2**, thanks @lshqqytiger!  
   *note*: new directml is finally based on modern `torch` 2.3.1!  
+- xyz grid: add support for LoRA selector
+- vae load: store original vae so it can be restored when set to none
 - extra networks: info display now contains link to source url if model if its known  
   works for civitai and huggingface models  
 - force gc for lowvram users and improve gc logging
@@ -55,13 +71,13 @@ But there's more than SD3:
 - additional torch gc checks, thanks @Disty0!
 
 **Improvements: NNCF**, thanks @Disty0!  
- - SD3 and PixArt support  
- - moved the first compression step to CPU  
- - sequential cpu offload (lowvram) support  
- - Lora support without reloading the model  
- - ControlNet compression support  
+- SD3 and PixArt support  
+- moved the first compression step to CPU  
+- sequential cpu offload (lowvram) support  
+- Lora support without reloading the model  
+- ControlNet compression support  
 
-### Fixes
+### Fixes for 2024-06-23
 
 - fix unsaturated outputs, force apply vae config on model load  
 - fix hidiffusion handling of non-square aspect ratios, thanks @ShenZhang-Shin!
@@ -79,6 +95,7 @@ But there's more than SD3:
 - fix api ip-adapter
 - fix memory exceptions with ROCm, thanks @Disty0!
 - fix face-hires with lowvram, thanks @Disty0!
+- fix pag incorrectly resetting pipeline
 - cleanup image metadata
 - restructure api examples: `cli/api-*`
 - handle theme fallback when invalid theme is specified
@@ -98,7 +115,7 @@ Plus tons of minor features such as optimized initial install experience, **T-Ga
 
 ### Full Changelog for 2024-06-13
 
-#### New Models
+#### New Models for 2024-06-23
 
 - [StabilityAI Stable Diffusion 3 Medium](https://stability.ai/news/stable-diffusion-3-medium)  
   yup, supported!  
@@ -109,7 +126,7 @@ Plus tons of minor features such as optimized initial install experience, **T-Ga
   note: this is a very large model at ~17GB, but can be used with less VRAM using model offloading  
   simply select from networks -> models -> reference, model will be auto-downloaded on first use  
 
-#### New Functionality
+#### New Functionality for 2024-06-23
 
 - [MuLan](https://github.com/mulanai/MuLan) Multi-language prompts
   write your prompts in ~110 auto-detected languages!  
@@ -146,7 +163,7 @@ Plus tons of minor features such as optimized initial install experience, **T-Ga
   typical differences are not large and its disabled by default as it does have some performance impact  
 - new sampler: **Euler FlowMatch**  
 
-#### Improvements
+#### Improvements Fixes 2024-06-13
 
 - additional modernui themes
 - reintroduce prompt attention normalization, disabled by default, enable in settings -> execution  
@@ -166,7 +183,7 @@ Plus tons of minor features such as optimized initial install experience, **T-Ga
 - auto-synchronize modernui and core branches  
 - add option to pad prompt with zeros, thanks @Disty
 
-#### Fixes
+#### Fixes 2024-06-13
 
 - cumulative fixes since the last release  
 - fix apply/unapply hidiffusion for sd15  

extensions-builtin/Lora/network_lora.py

@@ -26,7 +26,7 @@ class NetworkModuleLora(network.NetworkModule):
             return None
         linear_modules = [torch.nn.Linear, torch.nn.modules.linear.NonDynamicallyQuantizableLinear, torch.nn.MultiheadAttention, diffusers_lora.LoRACompatibleLinear]
         is_linear = type(self.sd_module) in linear_modules or self.sd_module.__class__.__name__ == "NNCFLinear"
-        is_conv = type(self.sd_module) in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv]
+        is_conv = type(self.sd_module) in [torch.nn.Conv2d, diffusers_lora.LoRACompatibleConv] or self.sd_module.__class__.__name__ == "NNCFConv2d"
         if is_linear:
             weight = weight.reshape(weight.shape[0], -1)
             module = torch.nn.Linear(weight.shape[1], weight.shape[0], bias=False)

installer.py

@@ -540,6 +540,27 @@ def install_rocm_zluda(torch_command):
             ort_version = os.environ.get('ONNXRUNTIME_VERSION', None)
             ort_package = os.environ.get('ONNXRUNTIME_PACKAGE', f"--pre onnxruntime-training{'' if ort_version is None else ('==' + ort_version)} --index-url https://pypi.lsh.sh/{rocm_ver[0]}{rocm_ver[2]} --extra-index-url https://pypi.org/simple")
             install(ort_package, 'onnxruntime-training')
+
+        if bool(int(os.environ.get("TORCH_BLAS_PREFER_HIPBLASLT", "1"))):
+            supported_archs = []
+            hipblaslt_available = True
+            libpath = os.environ.get("HIPBLASLT_TENSILE_LIBPATH", "/opt/rocm/lib/hipblaslt/library")
+            for file in os.listdir(libpath):
+                if not file.startswith('extop_'):
+                    continue
+                supported_archs.append(file[6:-3])
+            for gpu in amd_gpus:
+                if gpu not in supported_archs:
+                    hipblaslt_available = False
+                    break
+            log.info(f'hipBLASLt supported_archs={supported_archs}, available={hipblaslt_available}')
+            if hipblaslt_available:
+                import ctypes
+                # Preload hipBLASLt.
+                ctypes.CDLL("/opt/rocm/lib/libhipblaslt.so", mode=ctypes.RTLD_GLOBAL)
+                os.environ["HIPBLASLT_TENSILE_LIBPATH"] = libpath
+            else:
+                os.environ["TORCH_BLAS_PREFER_HIPBLASLT"] = "0"
     return torch_command
 
 
@@ -680,6 +701,20 @@ def check_torch():
         install('onnxruntime-gpu', 'onnxruntime-gpu', ignore=True, quiet=True)
     elif is_rocm_available(allow_rocm):
         torch_command = install_rocm_zluda(torch_command)
+
+        # WSL ROCm
+        if os.environ.get('WSL_DISTRO_NAME', None) is not None:
+            import ctypes
+            try:
+                # Preload stdc++ library. This will ignore Anaconda stdc++ library.
+                ctypes.CDLL("/lib/x86_64-linux-gnu/libstdc++.so.6", mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                pass
+            try:
+                # Preload HSA Runtime library.
+                ctypes.CDLL("/opt/rocm/lib/libhsa-runtime64.so", mode=ctypes.RTLD_GLOBAL)
+            except OSError:
+                log.error("Failed to preload HSA Runtime library.")
     elif is_ipex_available(allow_ipex):
         torch_command = install_ipex(torch_command)
     elif allow_openvino and args.use_openvino:
@@ -894,6 +929,7 @@ def install_submodules(force=True):
                 branch(name)
         except Exception:
             log.error(f'Error updating submodule: {submodule}')
+    setup_logging()
     if args.profile:
         print_profile(pr, 'Submodule')
     return '\n'.join(res)
@@ -1051,20 +1087,20 @@ def check_ui(ver):
 
     if not same(ver):
         log.debug(f'Branch mismatch: sdnext={ver["branch"]} ui={ver["ui"]}')
-    cwd = os.getcwd()
-    try:
-        os.chdir('extensions-builtin/sdnext-modernui')
-        target = 'dev' if 'dev' in ver['branch'] else 'main'
-        git('checkout ' + target, ignore=True, optional=True)
+        cwd = os.getcwd()
+        try:
+            os.chdir('extensions-builtin/sdnext-modernui')
+            target = 'dev' if 'dev' in ver['branch'] else 'main'
+            git('checkout ' + target, ignore=True, optional=True)
+            os.chdir(cwd)
+            ver = get_version(force=True)
+            if not same(ver):
+                log.debug(f'Branch synchronized: {ver["branch"]}')
+            else:
+                log.debug(f'Branch sync failed: sdnext={ver["branch"]} ui={ver["ui"]}')
+        except Exception as e:
+            log.debug(f'Branch switch: {e}')
         os.chdir(cwd)
-        ver = get_version(force=True)
-        if not same(ver):
-            log.debug(f'Branch synchronized: {ver["branch"]}')
-        else:
-            log.debug(f'Branch sync failed: sdnext={ver["branch"]} ui={ver["ui"]}')
-    except Exception as e:
-        log.debug(f'Branch switch: {e}')
-    os.chdir(cwd)
 
 
 # check version of the main repo and optionally upgrade it
@@ -1164,7 +1200,7 @@ def check_timestamp():
 def add_args(parser):
     group = parser.add_argument_group('Setup options')
     group.add_argument('--reset', default = os.environ.get("SD_RESET",False), action='store_true', help = "Reset main repository to latest version, default: %(default)s")
-    group.add_argument('--upgrade', default = os.environ.get("SD_UPGRADE",False), action='store_true', help = "Upgrade main repository to latest version, default: %(default)s")
+    group.add_argument('--upgrade', '--update', default = os.environ.get("SD_UPGRADE",False), action='store_true', help = "Upgrade main repository to latest version, default: %(default)s")
     group.add_argument('--requirements', default = os.environ.get("SD_REQUIREMENTS",False), action='store_true', help = "Force re-check of requirements, default: %(default)s")
     group.add_argument('--quick', default = os.environ.get("SD_QUICK",False), action='store_true', help = "Bypass version checks, default: %(default)s")
     group.add_argument('--use-directml', default = os.environ.get("SD_USEDIRECTML",False), action='store_true', help = "Use DirectML if no compatible GPU is detected, default: %(default)s")

javascript/extraNetworks.js

@@ -461,7 +461,7 @@ function setupExtraNetworksForTab(tabname) {
         en.style.position = 'absolute';
         en.style.right = '0';
         en.style.top = '13em';
-        en.style.height = '-webkit-fill-available';
+        en.style.height = 'auto';
         en.style.transition = 'width 0.3s ease';
         en.style.width = `${window.opts.extra_networks_sidebar_width}vw`;
         gradioApp().getElementById(`${tabname}_settings`).parentNode.style.width = `${100 - 2 - window.opts.extra_networks_sidebar_width}vw`;

modules/api/models.py

@@ -313,7 +313,7 @@ class ResInterrogate(BaseModel):
 
 class ReqVQA(BaseModel):
     image: str = Field(default="", title="Image", description="Image to work on, must be a Base64 string containing the image's data.")
-    model: str = Field(default="Moondream 2", title="Model", description="The interrogate model used.")
+    model: str = Field(default="MS Florence 2 Base", title="Model", description="The interrogate model used.")
     question: str = Field(default="describe the image", title="Question", description="Question to ask the model.")
 
 class ResVQA(BaseModel):

modules/api/script.py

@@ -39,10 +39,10 @@ def get_script(script_name, script_runner):
     return script_runner.scripts[script_idx]
 
 def init_default_script_args(script_runner):
-    #find max idx from the scripts in runner and generate a none array to init script_args
+    # find max idx from the scripts in runner and generate a none array to init script_args
     last_arg_index = 1
     for script in script_runner.scripts:
-        if last_arg_index < script.args_to:
+        if last_arg_index < script.args_to: # pylint disable=consider-using-max-builtin
             last_arg_index = script.args_to
     # None everywhere except position 0 to initialize script args
     script_args = [None]*last_arg_index

modules/control/run.py

@@ -351,6 +351,7 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini
     output_filename = None
     index = 0
     frames = 0
+    blended_image = None
 
     # set pipeline
     if pipe.__class__.__name__ != shared.sd_model.__class__.__name__:
@@ -477,7 +478,6 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini
                             process.model = None
 
                     debug(f'Control processed: {len(processed_images)}')
-                    blended_image = None
                     if len(processed_images) > 0:
                         try:
                             if len(p.extra_generation_params["Control process"]) == 0:
@@ -692,5 +692,4 @@ def control_run(units: List[unit.Unit] = [], inputs: List[Image.Image] = [], ini
     if is_generator:
         yield (output_images, blended_image, html_txt, output_filename)
     else:
-        yield (output_images, blended_image, html_txt, output_filename)
-        return
+        return (output_images, blended_image, html_txt, output_filename)

modules/control/units/controlnet.py

@@ -49,7 +49,8 @@ predefined_sdxl = {
     'Canny XL': 'diffusers/controlnet-canny-sdxl-1.0',
     'Depth Zoe XL': 'diffusers/controlnet-zoe-depth-sdxl-1.0',
     'Depth Mid XL': 'diffusers/controlnet-depth-sdxl-1.0-mid',
-    'OpenPose XL': 'thibaud/controlnet-openpose-sdxl-1.0',
+    'OpenPose XL': 'thibaud/controlnet-openpose-sdxl-1.0/bin',
+    # 'OpenPose XL': 'thibaud/controlnet-openpose-sdxl-1.0/OpenPoseXL2.safetensors',
     'Xinsir OpenPose XL': 'xinsir/controlnet-openpose-sdxl-1.0',
     'Xinsir Canny XL': 'xinsir/controlnet-canny-sdxl-1.0',
     'Xinsir Scribble XL': 'xinsir/controlnet-scribble-sdxl-1.0',
@@ -171,6 +172,9 @@ class ControlNet():
             if model_path.endswith('.safetensors'):
                 self.load_safetensors(model_path)
             else:
+                if '/bin' in model_path:
+                    model_path = model_path.replace('/bin', '')
+                    self.load_config['use_safetensors'] = False
                 self.model = ControlNetModel.from_pretrained(model_path, **self.load_config)
             if self.dtype is not None:
                 self.model.to(self.dtype)

modules/devices.py

@@ -46,7 +46,7 @@ def get_gpu_info():
         try:
             if shared.cmd_opts.use_openvino:
                 return {
-                    'device': get_openvino_device(),
+                    'device': get_openvino_device(), # pylint: disable=used-before-assignment
                     'openvino': get_package_version("openvino"),
                 }
             elif shared.cmd_opts.use_directml:
@@ -311,7 +311,7 @@ def set_cuda_params():
         inference_context = contextlib.nullcontext
     else:
         inference_context = torch.no_grad
-    log_device_name = get_raw_openvino_device() if shared.cmd_opts.use_openvino else torch.device(get_optimal_device_name())
+    log_device_name = get_raw_openvino_device() if shared.cmd_opts.use_openvino else torch.device(get_optimal_device_name()) # pylint: disable=used-before-assignment
     log.debug(f'Desired Torch parameters: dtype={shared.opts.cuda_dtype} no-half={shared.opts.no_half} no-half-vae={shared.opts.no_half_vae} upscast={shared.opts.upcast_sampling}')
     log.info(f'Setting Torch parameters: device={log_device_name} dtype={dtype} vae={dtype_vae} unet={dtype_unet} context={inference_context.__name__} fp16={fp16_ok} bf16={bf16_ok} optimization={shared.opts.cross_attention_optimization}')
 

modules/loader.py

@@ -35,7 +35,7 @@ timer.startup.record("torch")
 import transformers # pylint: disable=W0611,C0411
 timer.startup.record("transformers")
 
-import onnxruntime
+import onnxruntime # pylint: disable=W0611,C0411
 onnxruntime.set_default_logger_severity(3)
 timer.startup.record("onnx")
 
@@ -50,7 +50,7 @@ timer.startup.record("pydantic")
 import diffusers # pylint: disable=W0611,C0411
 import diffusers.loaders.single_file # pylint: disable=W0611,C0411
 logging.getLogger("diffusers.loaders.single_file").setLevel(logging.ERROR)
-from tqdm.rich import tqdm
+from tqdm.rich import tqdm # pylint: disable=W0611,C0411
 diffusers.loaders.single_file.logging.tqdm = partial(tqdm, unit='C')
 timer.startup.record("diffusers")
 

modules/model_sd3.py

@@ -13,9 +13,9 @@ def load_sd3(fn=None, cache_dir=None, config=None):
     if fn is not None and fn.endswith('.safetensors') and os.path.exists(fn):
         model_id = fn
         loader = diffusers.StableDiffusion3Pipeline.from_single_file
-        diffusers_minor = int(diffusers.__version__.split('.')[1])
+        _diffusers_major, diffusers_minor, diffusers_micro = int(diffusers.__version__.split('.')[0]), int(diffusers.__version__.split('.')[1]), int(diffusers.__version__.split('.')[2])
         fn_size = os.path.getsize(fn)
-        if diffusers_minor < 30 or fn_size < 5e9: # te1/te2 do not get loaded correctly in diffusers 0.29.0 or model is without te1/te2
+        if (diffusers_minor <= 29 and diffusers_micro < 1) or fn_size < 5e9: # te1/te2 do not get loaded correctly in diffusers 0.29.0 if model is without te1/te2
             kwargs = {
                 'text_encoder': transformers.CLIPTextModelWithProjection.from_pretrained(
                     repo_id,

modules/model_t5.py

@@ -75,3 +75,4 @@ def set_t5(pipe, module, t5=None, cache_dir=None):
         else:
             pipe.maybe_free_model_hooks()
     devices.torch_gc()
+    return pipe

modules/pag/__init__.py

@@ -15,9 +15,11 @@ def apply(p: processing.StableDiffusionProcessing): # pylint: disable=arguments-
     c = shared.sd_model.__class__ if shared.sd_loaded else None
     if c == StableDiffusionPAGPipeline or c == StableDiffusionXLPAGPipeline:
         unapply()
-        return None
     if p.pag_scale == 0:
         return
+    if sd_models.get_diffusers_task(shared.sd_model) != sd_models.DiffusersTaskType.TEXT_2_IMAGE:
+        shared.log.warning(f'PAG: pipeline={c} not implemented')
+        return None
     if detect.is_sd15(c):
         orig_pipeline = shared.sd_model
         shared.sd_model = sd_models.switch_pipe(StableDiffusionPAGPipeline, shared.sd_model)

modules/processing.py

@@ -158,7 +158,6 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
         shared.prompt_styles.apply_styles_to_extra(p)
         shared.prompt_styles.extract_comments(p)
-        pag.apply(p)
         if shared.opts.cuda_compile_backend == 'none':
             sd_models.apply_token_merging(p.sd_model)
             sd_hijack_freeu.apply_freeu(p, not shared.native)
@@ -273,6 +272,7 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
         extra_network_data = None
         debug(f'Processing inner: args={vars(p)}')
         for n in range(p.n_iter):
+            pag.apply(p)
             debug(f'Processing inner: iteration={n+1}/{p.n_iter}')
             p.iteration = n
             if shared.state.skipped:

modules/processing_args.py

@@ -27,6 +27,8 @@ def task_specific_kwargs(p, model):
                 'height': 8 * math.ceil(p.height / 8),
             }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
+        if shared.sd_model_type == 'sdxl':
+            model.register_to_config(requires_aesthetics_score = False)
         p.ops.append('img2img')
         task_args = {
             'image': p.init_images,
@@ -41,6 +43,8 @@ def task_specific_kwargs(p, model):
             'strength': p.denoising_strength,
         }
     elif (sd_models.get_diffusers_task(model) == sd_models.DiffusersTaskType.INPAINTING or is_img2img_model) and len(getattr(p, 'init_images', [])) > 0:
+        if shared.sd_model_type == 'sdxl':
+            model.register_to_config(requires_aesthetics_score = False)
         p.ops.append('inpaint')
         width, height = processing_helpers.resize_init_images(p)
         task_args = {

modules/processing_callbacks.py

@@ -38,7 +38,8 @@ def diffusers_callback(pipe, step: int, timestep: int, kwargs: dict):
         return kwargs
     latents = kwargs.get('latents', None)
     debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}')
-    shared.state.sampling_step = step
+    order = getattr(pipe.scheduler, "order", 1) if hasattr(pipe, 'scheduler') else 1
+    shared.state.sampling_step = step // order
     if shared.state.interrupted or shared.state.skipped:
         raise AssertionError('Interrupted...')
     if shared.state.paused:

modules/processing_diffusers.py

@@ -85,7 +85,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
 
     shared.sd_model = update_pipeline(shared.sd_model, p)
     shared.log.info(f'Base: class={shared.sd_model.__class__.__name__}')
-    update_sampler(p, shared.sd_model) # TODO SD3
+    update_sampler(p, shared.sd_model)
     base_args = set_pipeline_args(
         p=p,
         model=shared.sd_model,

modules/prompt_parser_diffusers.py

@@ -168,6 +168,11 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
         p.negative_embeds = []
         p.negative_pooleds = []
 
+        if (shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload) and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"):
+            # if the last job is interrupted, model will stay in the vram and cause oom, send everything back to cpu before continuing
+            pipe.maybe_free_model_hooks()
+            devices.torch_gc()
+
         for i in range(max(len(positive_schedule), len(negative_schedule))):
             positive_prompt = positive_schedule[i % len(positive_schedule)]
             negative_prompt = negative_schedule[i % len(negative_schedule)]
@@ -199,7 +204,11 @@ def encode_prompts(pipe, p, prompts: list, negative_prompts: list, steps: int, c
         if debug_enabled:
             get_tokens('positive', prompts[0])
             get_tokens('negative', negative_prompts[0])
+        if (shared.cmd_opts.medvram or shared.opts.diffusers_model_cpu_offload) and hasattr(pipe, "_all_hooks") and hasattr(pipe, "maybe_free_model_hooks"):
+            # text encoder will stay in the vram and cause oom, send everything back to cpu before continuing
+            pipe.maybe_free_model_hooks()
         debug(f"Prompt encode: time={(time.time() - t0):.3f}")
+        devices.torch_gc()
         return
 
 

modules/scripts.py

@@ -489,10 +489,9 @@ class ScriptRunner:
         s = ScriptSummary('before-process')
         for script in self.alwayson_scripts:
             try:
-                args = p.script_args[script.args_from:script.args_to]
-                if len(args) == 0:
-                    continue
-                script.before_process(p, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.before_process(p, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f"Error running before process: {script.filename}")
             s.record(script.title())
@@ -502,10 +501,9 @@ class ScriptRunner:
         s = ScriptSummary('process')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.process(p, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.process(p, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script process: {script.filename}')
             s.record(script.title())
@@ -516,10 +514,9 @@ class ScriptRunner:
         processed = None
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                processed = script.process_images(p, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    processed = script.process_images(p, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script process images: {script.filename}')
             s.record(script.title())
@@ -530,10 +527,9 @@ class ScriptRunner:
         s = ScriptSummary('before-process-batch')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.before_process_batch(p, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.before_process_batch(p, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script before process batch: {script.filename}')
             s.record(script.title())
@@ -543,10 +539,9 @@ class ScriptRunner:
         s = ScriptSummary('process-batch')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.process_batch(p, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.process_batch(p, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script process batch: {script.filename}')
             s.record(script.title())
@@ -556,10 +551,9 @@ class ScriptRunner:
         s = ScriptSummary('postprocess')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.postprocess(p, processed, *args)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.postprocess(p, processed, *args)
             except Exception as e:
                 errors.display(e, f'Running script postprocess: {script.filename}')
             s.record(script.title())
@@ -569,10 +563,9 @@ class ScriptRunner:
         s = ScriptSummary('postprocess-batch')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.postprocess_batch(p, *args, images=images, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.postprocess_batch(p, *args, images=images, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script before postprocess batch: {script.filename}')
             s.record(script.title())
@@ -582,10 +575,9 @@ class ScriptRunner:
         s = ScriptSummary('postprocess-batch-list')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.postprocess_batch_list(p, pp, *args, **kwargs)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.postprocess_batch_list(p, pp, *args, **kwargs)
             except Exception as e:
                 errors.display(e, f'Running script before postprocess batch list: {script.filename}')
             s.record(script.title())
@@ -595,10 +587,9 @@ class ScriptRunner:
         s = ScriptSummary('postprocess-image')
         for script in self.alwayson_scripts:
             try:
-                args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
-                if len(args) == 0:
-                    continue
-                script.postprocess_image(p, pp, *args)
+                if (script.args_to > 0) and (script.args_to >= script.args_from):
+                    args = p.per_script_args.get(script.title(), p.script_args[script.args_from:script.args_to])
+                    script.postprocess_image(p, pp, *args)
             except Exception as e:
                 errors.display(e, f'Running script postprocess image: {script.filename}')
             s.record(script.title())

modules/sd_models.py

@@ -38,6 +38,7 @@ sd_metadata_pending = 0
 sd_metadata_timer = 0
 debug_move = shared.log.trace if os.environ.get('SD_MOVE_DEBUG', None) is not None else lambda *args, **kwargs: None
 debug_load = os.environ.get('SD_LOAD_DEBUG', None)
+debug_process = shared.log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
 diffusers_version = int(diffusers.__version__.split('.')[1])
 
 
@@ -546,7 +547,7 @@ def change_backend():
     shared.native = shared.backend == shared.Backend.DIFFUSERS
     checkpoints_loaded.clear()
     from modules.sd_samplers import list_samplers
-    list_samplers(shared.backend)
+    list_samplers()
     list_models()
     from modules.sd_vae import refresh_vae_list
     refresh_vae_list()
@@ -585,7 +586,7 @@ def detect_pipeline(f: str, op: str = 'model', warning=True, quiet=False):
                     guess = 'Stable Diffusion XL Instruct'
                 elif (size > 3138 and size < 3142): #3140
                     guess = 'Stable Diffusion XL'
-                elif (size > 5692 and size < 5698) or (size > 4134 and size < 4138) or (size > 10362 and size < 10366):
+                elif (size > 5692 and size < 5698) or (size > 4134 and size < 4138) or (size > 10362 and size < 10366) or (size > 15028 and size < 15228):
                     guess = 'Stable Diffusion 3'
             # guess by name
             """
@@ -1297,6 +1298,7 @@ def switch_pipe(cls: diffusers.DiffusionPipeline, pipeline: diffusers.DiffusionP
 
 def clean_diffuser_pipe(pipe):
     if pipe is not None and shared.sd_model_type == 'sdxl' and 'requires_aesthetics_score' in pipe.config and hasattr(pipe, '_internal_dict'):
+        debug_process(f'Pipeline clean: {pipe.__class__.__name__}')
         # diffusers adds requires_aesthetics_score with img2img and complains if requires_aesthetics_score exist in txt2img
         internal_dict = dict(pipe._internal_dict) # pylint: disable=protected-access
         internal_dict.pop('requires_aesthetics_score', None)

modules/sd_models_compile.py

@@ -309,6 +309,9 @@ def check_deepcache(enable: bool):
 
 def compile_deepcache(sd_model):
     global deepcache_worker # pylint: disable=global-statement
+    if not hasattr(sd_model, 'unet'):
+        shared.log.warning(f'Model compile using deep-cache: {sd_model.__class__} not supported')
+        return sd_model
     try:
         from DeepCache import DeepCacheSDHelper
     except Exception as e:

modules/sd_samplers.py

@@ -14,7 +14,7 @@ samplers_map = {}
 loaded_config = None
 
 
-def list_samplers(backend_name = shared.backend):
+def list_samplers():
     global all_samplers # pylint: disable=global-statement
     global all_samplers_map # pylint: disable=global-statement
     global samplers # pylint: disable=global-statement

modules/sd_samplers_cfg_denoiser.py

@@ -1,5 +1,5 @@
-# TODO a1111 compatibility module
-# TODO cfg_denoiser implementation missing
+# a1111 compatibility module
+# cfg_denoiser implementation missing
 
 import torch
 from modules import prompt_parser, devices, sd_samplers_common
@@ -95,7 +95,7 @@ class CFGDenoiser(torch.nn.Module):
         if state.interrupted or state.skipped:
             raise sd_samplers_common.InterruptedException
 
-        # TODO cfg_scale implementation missing
+        # cfg_scale implementation missing for original backend
         # if sd_samplers_common.apply_refiner(self):
         #     cond = self.sampler.sampler_extra_args['cond']
         #    uncond = self.sampler.sampler_extra_args['uncond']

modules/sd_samplers_timesteps.py

@@ -1,4 +1,4 @@
-# TODO a1111 compatibility module
+# a1111 compatibility module
 
 import torch
 from modules import sd_samplers_common, sd_samplers_timesteps_impl, sd_samplers_compvis

modules/sd_samplers_timesteps_impl.py

@@ -1,4 +1,4 @@
-# TODO a1111 compatibility module
+# a1111 compatibility module
 
 import torch
 import tqdm

modules/sd_vae.py

@@ -259,6 +259,11 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         vae_file, vae_source = resolve_vae(checkpoint_file)
     else:
         vae_source = "function-argument"
+    if vae_file is None or vae_file == 'None':
+        if hasattr(sd_model, 'original_vae'):
+            sd_models.set_diffuser_options(sd_model, vae=sd_model.original_vae, op='vae')
+            shared.log.info("VAE restored")
+            return None
     if loaded_vae_file == vae_file:
         return None
     if not shared.native and (shared.cmd_opts.lowvram or shared.cmd_opts.medvram):
@@ -276,11 +281,14 @@ def reload_vae_weights(sd_model=None, vae_file=unspecified):
         if vae_file is not None:
             shared.log.info(f"VAE weights loaded: {vae_file}")
     else:
-        if hasattr(shared.sd_model, "vae") and hasattr(shared.sd_model, "sd_checkpoint_info"):
-            vae = load_vae_diffusers(shared.sd_model.sd_checkpoint_info.filename, vae_file, vae_source)
+        if hasattr(sd_model, "vae") and hasattr(sd_model, "sd_checkpoint_info"):
+            vae = load_vae_diffusers(sd_model.sd_checkpoint_info.filename, vae_file, vae_source)
             if vae is not None:
+                if not hasattr(sd_model, 'original_vae'):
+                    sd_model.original_vae = sd_model.vae
+                    sd_models.move_model(sd_model.original_vae, devices.cpu)
                 sd_models.set_diffuser_options(sd_model, vae=vae, op='vae')
-                apply_vae_config(shared.sd_model.sd_checkpoint_info.filename, vae_file, sd_model)
+                apply_vae_config(sd_model.sd_checkpoint_info.filename, vae_file, sd_model)
 
     if not shared.cmd_opts.lowvram and not shared.cmd_opts.medvram:
         sd_models.move_model(sd_model, devices.device)

modules/shared.py

@@ -813,7 +813,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
     "extra_networks_sort": OptionInfo("Default", "Sort order", gr.Dropdown, {"choices": ['Default', 'Name [A-Z]', 'Name [Z-A]', 'Date [Newest]', 'Date [Oldest]', 'Size [Largest]', 'Size [Smallest]']}),
     "extra_networks_view": OptionInfo("gallery", "UI view", gr.Radio, {"choices": ["gallery", "list"]}),
     "extra_networks_card_cover": OptionInfo("sidebar", "UI position", gr.Radio, {"choices": ["cover", "inline", "sidebar"]}),
-    "extra_networks_height": OptionInfo(53, "UI height (%)", gr.Slider, {"minimum": 10, "maximum": 100, "step": 1}),
+    "extra_networks_height": OptionInfo(55, "UI height (%)", gr.Slider, {"minimum": 10, "maximum": 100, "step": 1}),
     "extra_networks_sidebar_width": OptionInfo(35, "UI sidebar width (%)", gr.Slider, {"minimum": 10, "maximum": 80, "step": 1}),
     "extra_networks_card_size": OptionInfo(160, "UI card size (px)", gr.Slider, {"minimum": 20, "maximum": 2000, "step": 1}),
     "extra_networks_card_square": OptionInfo(True, "UI disable variable aspect ratio"),
@@ -821,7 +821,7 @@ options_templates.update(options_section(('extra_networks', "Networks"), {
     "extra_networks_sep2": OptionInfo("<h2>Extra networks general</h2>", "", gr.HTML),
     "extra_network_reference": OptionInfo(False, "Use reference values when available", gr.Checkbox),
     "extra_network_skip_indexing": OptionInfo(False, "Build info on first access", gr.Checkbox),
-    "extra_networks_default_multiplier": OptionInfo(1.0, "Default multiplier for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
+    "extra_networks_default_multiplier": OptionInfo(1.0, "Default strength for extra networks", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     "diffusers_convert_embed": OptionInfo(False, "Auto-convert SD 1.5 embeddings to SDXL ", gr.Checkbox, {"visible": native}),
     "extra_networks_sep3": OptionInfo("<h2>Extra networks settings</h2>", "", gr.HTML),
     "extra_networks_styles": OptionInfo(True, "Show built-in styles"),

modules/textual_inversion/textual_inversion.py

@@ -311,14 +311,14 @@ class EmbeddingDatabase:
 
         if ext in ['.PNG', '.WEBP', '.JXL', '.AVIF']:
             if '.preview' in filename.lower():
-                return
+                return None
             embed_image = Image.open(path)
             if hasattr(embed_image, 'text') and 'sd-ti-embedding' in embed_image.text:
                 data = embedding_from_b64(embed_image.text['sd-ti-embedding'])
             else:
                 data = extract_image_data_embed(embed_image)
                 if not data: # if data is None, means this is not an embeding, just a preview image
-                    return
+                    return None
         elif ext in ['.BIN', '.PT']:
             data = torch.load(path, map_location="cpu")
         elif ext in ['.SAFETENSORS']:
@@ -336,7 +336,7 @@ class EmbeddingDatabase:
         elif type(data) == dict and type(next(iter(data.values()))) == torch.Tensor:
             if len(data.keys()) != 1:
                 self.skipped_embeddings[name] = Embedding(None, name=name, filename=path)
-                return
+                return None
             emb = next(iter(data.values()))
             if len(emb.shape) == 1:
                 emb = emb.unsqueeze(0)

modules/ui_extensions.py

@@ -206,7 +206,7 @@ def uninstall_extension(extension_path, search_text, sort_column):
     if len(found) > 0 and os.path.isdir(extension_path):
         found = found[0]
         try:
-            shutil.rmtree(found.path, ignore_errors=False, onerror=errorRemoveReadonly)
+            shutil.rmtree(found.path, ignore_errors=False, onerror=errorRemoveReadonly) # pylint: disable=deprecated-argument
             # extensions.extensions = [extension for extension in extensions.extensions if os.path.abspath(found.path) != os.path.abspath(extension_path)]
         except Exception as e:
             shared.log.warning(f'Extension uninstall failed: {found.path} {e}')

modules/ui_postprocessing.py

@@ -80,7 +80,7 @@ def create_ui():
                     with gr.Row():
                         vqa_answer = gr.Textbox(label="Answer", lines=3)
                     with gr.Row(elem_id='interrogate_buttons_query'):
-                        vqa_model = gr.Dropdown(list(vqa.MODELS), value='Moondream 2', label='VQA Model')
+                        vqa_model = gr.Dropdown(list(vqa.MODELS), value='MS Florence 2 Base', label='VQA Model')
                         vqa_submit = gr.Button("Interrogate", elem_id="interrogate_btn_interrogate", variant='primary')
                     vqa_submit.click(vqa.interrogate, inputs=[vqa_question, vqa_image, vqa_model], outputs=[vqa_answer])
 

modules/vqa.py

@@ -1,5 +1,6 @@
 import torch
 import transformers
+import transformers.dynamic_module_utils
 from PIL import Image
 from modules import shared, devices
 
@@ -8,6 +9,8 @@ processor = None
 model = None
 loaded: str = None
 MODELS = {
+    "MS Florence 2 Base": "microsoft/Florence-2-base", # 0.5GB
+    "MS Florence 2 Large": "microsoft/Florence-2-large", # 1.5GB
     "Moondream 2": "vikhyatk/moondream2", # 3.7GB
     "GIT TextCaps Base": "microsoft/git-base-textcaps", # 0.7GB
     "GIT VQA Base": "microsoft/git-base-vqav2", # 0.7GB
@@ -124,6 +127,49 @@ def moondream(question: str, image: Image.Image, repo: str = None):
     return response
 
 
+def florence(question: str, image: Image.Image, repo: str = None):
+    global processor, model, loaded # pylint: disable=global-statement
+    _get_imports = transformers.dynamic_module_utils.get_imports
+    def get_imports(f):
+        R = _get_imports(f)
+        if "flash_attn" in R:
+            R.remove("flash_attn") # flash_attn is optional
+        return R
+    if model is None or loaded != repo:
+        transformers.dynamic_module_utils.get_imports = get_imports
+        model = transformers.AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+        processor = transformers.AutoProcessor.from_pretrained(repo, trust_remote_code=True)
+        transformers.dynamic_module_utils.get_imports = _get_imports
+        loaded = repo
+        model.eval()
+    model.to(devices.device, devices.dtype)
+    shared.log.debug(f'VQA: class={model.__class__.__name__} processor={processor.__class__} model={repo}')
+
+    if question.startswith('<'):
+        task = question.split('>', 1)[0] + '>'
+    else:
+        task = '<MORE_DETAILED_CAPTION>'
+        question = task + question
+    inputs = processor(text=question, images=image, return_tensors="pt")
+    input_ids = inputs['input_ids'].to(devices.device)
+    pixel_values = inputs['pixel_values'].to(devices.device, devices.dtype)
+    with devices.inference_context():
+        generated_ids = model.generate(
+            input_ids=input_ids,
+            pixel_values=pixel_values,
+            max_new_tokens=1024,
+            num_beams=3,
+            do_sample=False
+        )
+        generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
+        response = processor.post_process_generation(generated_text, task="task", image_size=(image.width, image.height))
+
+    if 'task' in response:
+        response = response['task']
+    shared.log.debug(f'VQA: task={task} response="{response}"')
+    return response
+
+
 def interrogate(vqa_question, vqa_image, vqa_model_req):
     vqa_model = MODELS.get(vqa_model_req, None)
     shared.log.debug(f'VQA: model="{vqa_model}" question="{vqa_question}" image={vqa_image}')
@@ -138,14 +184,16 @@ def interrogate(vqa_question, vqa_image, vqa_model_req):
         return answer
     if 'git' in vqa_model.lower():
         answer = git(vqa_question, vqa_image, vqa_model)
-    if 'vilt' in vqa_model.lower():
+    elif 'vilt' in vqa_model.lower():
         answer = vilt(vqa_question, vqa_image, vqa_model)
-    if 'blip' in vqa_model.lower():
+    elif 'blip' in vqa_model.lower():
         answer = blip(vqa_question, vqa_image, vqa_model)
-    if 'pix' in vqa_model.lower():
+    elif 'pix' in vqa_model.lower():
         answer = pix(vqa_question, vqa_image, vqa_model)
-    if 'moondream2' in vqa_model.lower():
+    elif 'moondream2' in vqa_model.lower():
         answer = moondream(vqa_question, vqa_image, vqa_model)
+    elif 'florence' in vqa_model.lower():
+        answer = florence(vqa_question, vqa_image, vqa_model)
     else:
         answer = 'unknown model'
     if model is not None:

modules/zluda.py

@@ -1,4 +1,3 @@
-import os
 import sys
 from typing import Union
 import torch
@@ -6,17 +5,13 @@ from torch._prims_common import DeviceLikeType
 import onnxruntime as ort
 from modules import shared, devices
 from modules.onnx_impl.execution_providers import available_execution_providers, ExecutionProvider
+from modules.zluda_hijacks import do_hijack
 
 
 PLATFORM = sys.platform
 do_nothing = lambda _: None # pylint: disable=unnecessary-lambda-assignment
 
 
-def _join_rocm_home(*paths) -> str:
-    from torch.utils.cpp_extension import ROCM_HOME
-    return os.path.join(ROCM_HOME, *paths)
-
-
 def is_zluda(device: DeviceLikeType):
     try:
         device = torch.device(device)
@@ -42,16 +37,9 @@ def initialize_zluda():
     if not devices.cuda_ok or not is_zluda(device):
         return
 
-    torch.version.hip = "5.7"
-    sys.platform = ""
-    from torch.utils import cpp_extension
-    sys.platform = PLATFORM
-    cpp_extension.IS_WINDOWS = PLATFORM == "win32"
-    cpp_extension.IS_MACOS = False
-    cpp_extension.IS_LINUX = sys.platform.startswith('linux')
-    cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access
+    do_hijack()
 
-    if cpp_extension.IS_WINDOWS:
+    if PLATFORM == "win32":
         torch.backends.cudnn.enabled = False
         torch.backends.cuda.enable_flash_sdp(False)
         torch.backends.cuda.enable_flash_sdp = do_nothing

modules/zluda_hijacks.py

@@ -0,0 +1,28 @@
+import os
+import sys
+import torch
+
+
+_topk = torch.topk
+def topk(tensor: torch.Tensor, *args, **kwargs):
+    device = tensor.device
+    values, indices = _topk(tensor.cpu(), *args, **kwargs)
+    return torch.return_types.topk((values.to(device), indices.to(device),))
+
+
+def _join_rocm_home(*paths) -> str:
+    from torch.utils.cpp_extension import ROCM_HOME
+    return os.path.join(ROCM_HOME, *paths)
+
+
+def do_hijack():
+    torch.version.hip = "5.7"
+    torch.topk = topk
+    platform = sys.platform
+    sys.platform = ""
+    from torch.utils import cpp_extension
+    sys.platform = platform
+    cpp_extension.IS_WINDOWS = platform == "win32"
+    cpp_extension.IS_MACOS = False
+    cpp_extension.IS_LINUX = platform.startswith('linux')
+    cpp_extension._join_rocm_home = _join_rocm_home # pylint: disable=protected-access

modules/zluda_installer.py

@@ -33,7 +33,7 @@ def install(zluda_path: os.PathLike) -> None:
     if os.path.exists(zluda_path):
         return
 
-    if platform.system() != 'Windows': # TODO
+    if platform.system() != 'Windows': # Windows-only. (PyTorch should be rebuilt on Linux)
         return
 
     urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/{RELEASE}/ZLUDA-windows-amd64.zip', '_zluda')

scripts/xyz_grid.py

@@ -65,6 +65,7 @@ def apply_sampler(p, x, xs):
     else:
         p.sampler_name = sampler_name
 
+
 def apply_hr_sampler_name(p, x, xs):
     hr_sampler_name = sd_samplers.samplers_map.get(x.lower(), None)
     if hr_sampler_name is None:
@@ -72,6 +73,7 @@ def apply_hr_sampler_name(p, x, xs):
     else:
         p.hr_sampler_name = hr_sampler_name
 
+
 def confirm_samplers(p, xs):
     for x in xs:
         if x.lower() not in sd_samplers.samplers_map:
@@ -138,6 +140,19 @@ def apply_vae(p, x, xs):
     sd_vae.reload_vae_weights(shared.sd_model, vae_file=find_vae(x))
 
 
+def list_lora():
+    import sys
+    lora = [v for k, v in sys.modules.items() if k == 'networks'][0]
+    loras = [v.name for v in lora.available_networks.values()]
+    return ['None'] + loras
+
+
+def apply_lora(p, x, xs):
+    if x == 'None':
+        return
+    p.prompt = p.prompt + f" <lora:{x}:{shared.opts.extra_networks_default_multiplier}>"
+
+
 def apply_te(p, x, xs):
     shared.opts.data["sd_text_encoder"] = x
     sd_models.reload_text_encoder()
@@ -235,6 +250,7 @@ axis_options = [
     AxisOption("Prompt S/R", str, apply_prompt, fmt=format_value),
     AxisOption("Model", str, apply_checkpoint, fmt=format_value, cost=1.0, choices=lambda: sorted(sd_models.checkpoints_list)),
     AxisOption("VAE", str, apply_vae, cost=0.7, choices=lambda: ['None'] + list(sd_vae.vae_dict)),
+    AxisOption("LoRA", str, apply_lora, cost=0.5, choices=list_lora),
     AxisOption("Text encoder", str, apply_te, cost=0.7, choices=lambda: ['None', 'T5 FP4', 'T5 FP8', 'T5 FP16']),
     AxisOption("Styles", str, apply_styles, choices=lambda: [s.name for s in shared.prompt_styles.styles.values()]),
     AxisOption("Seed", int, apply_field("seed")),