Merge branch 'dev' into xhinker

modules/devices.py

@@ -3,8 +3,8 @@ import gc
 import sys
 import time
 import contextlib
-import torch
 from functools import wraps
+import torch
 from modules.errors import log
 from modules import cmd_args, shared, memstats, errors
 

modules/model_flux.py

@@ -1,17 +1,16 @@
+import os
 import json
 import torch
 import diffusers
 import transformers
 from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download
-
+from accelerate.utils import compute_module_sizes
 from modules import shared, devices
 
 
-
 def load_quanto_transformer(repo_path):
-    from optimum.quanto import requantize
-    with open(repo_path + "/" + "transformer/quantization_map.json", "r") as f:
+    from optimum.quanto import requantize # pylint: disable=no-name-in-module
+    with open(repo_path + "/" + "transformer/quantization_map.json", "r", encoding='utf8') as f:
         quantization_map = json.load(f)
     state_dict = load_file(repo_path + "/" + "transformer/diffusion_pytorch_model.safetensors")
     dtype = state_dict['context_embedder.bias'].dtype
@@ -23,10 +22,10 @@ def load_quanto_transformer(repo_path):
 
 
 def load_quanto_text_encoder_2(repo_path):
-    from optimum.quanto import requantize
-    with open(repo_path + "/" + "text_encoder_2/quantization_map.json", "r") as f:
+    from optimum.quanto import requantize # pylint: disable=no-name-in-module
+    with open(repo_path + "/" + "text_encoder_2/quantization_map.json", "r", encoding='utf8') as f:
         quantization_map = json.load(f)
-    with open(repo_path + "/" + "text_encoder_2/config.json") as f:
+    with open(repo_path + "/" + "text_encoder_2/config.json", encoding='utf8') as f:
         t5_config = transformers.T5Config(**json.load(f))
     state_dict = load_file(repo_path + "/" + "text_encoder_2/model.safetensors")
     dtype = state_dict['encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight'].dtype
@@ -36,12 +35,35 @@ def load_quanto_text_encoder_2(repo_path):
     text_encoder_2.eval()
     return text_encoder_2
 
+
 def load_flux(checkpoint_info, diffusers_load_config):
-    if "qint8" in checkpoint_info.name.lower() or "qint4" in checkpoint_info.name.lower():
-        shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant=True')
+    if "qint8" in checkpoint_info.path.lower():
+        quant = 'qint8'
+    elif "qint4" in checkpoint_info.path.lower():
+        quant = 'qint4'
+    elif "nf4" in checkpoint_info.path.lower():
+        quant = 'nf4'
+    else:
+        quant = None
+    shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant={quant}')
+    if quant == 'nf4':
+        from installer import install
+        install('bitsandbytes', quiet=True)
+        try:
+            import bitsandbytes # pylint: disable=unused-import
+        except Exception as e:
+            shared.log.error(f"FLUX: Failed to import bitsandbytes: {e}")
+            raise
+        from modules.model_flux_nf4 import load_flux_nf4
+        pipe = load_flux_nf4(checkpoint_info, diffusers_load_config)
+    elif quant == 'qint8' or quant == 'qint4':
         from installer import install
         install('optimum-quanto', quiet=True)
-        from optimum import quanto
+        try:
+            from optimum import quanto # pylint: disable=no-name-in-module
+        except Exception as e:
+            shared.log.error(f"FLUX: Failed to import optimum-quanto: {e}")
+            raise
         quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
         pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, transformer=None, text_encoder_2=None, **diffusers_load_config)
         pipe.transformer = load_quanto_transformer(checkpoint_info.path)
@@ -50,18 +72,18 @@ def load_flux(checkpoint_info, diffusers_load_config):
             try:
                 pipe.transformer = pipe.transformer.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"FLUX: Failed to cast the transformer to {devices.dtype}! Set dtype to {pipe.transformer.dtype}")
+                shared.log.error(f"FLUX: Failed to cast transformer to {devices.dtype}, set dtype to {pipe.transformer.dtype}")
                 raise
         if pipe.text_encoder_2.dtype != devices.dtype:
             try:
                 pipe.text_encoder_2 = pipe.text_encoder_2.to(dtype=devices.dtype)
             except Exception:
-                shared.log.error(f"FLUX: Failed to cast the text encoder to {devices.dtype}! Set dtype to {pipe.text_encoder_2.dtype}")
+                shared.log.error(f"FLUX: Failed to cast text encoder to {devices.dtype}, set dtype to {pipe.text_encoder_2.dtype}")
                 raise
     else:
         pipe = diffusers.FluxPipeline.from_pretrained(checkpoint_info.path, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
-        shared.log.debug(f'Loading FLUX: model="{checkpoint_info.name}" quant=False')
     if devices.dtype == torch.float16 and not shared.opts.no_half_vae:
-        shared.log.warning("FLUX VAE doesn't support FP16! Enabling no-half-vae")
+        shared.log.warning("FLUX: does not support FP16 VAE, enabling no-half-vae")
         shared.opts.no_half_vae = True
+    shared.log.debug(f'FLUX computed size: {round(compute_module_sizes(pipe.transformer)[""] / 1024 / 1204)}')
     return pipe

modules/model_flux_nf4.py

@@ -0,0 +1,190 @@
+"""
+Copied from: https://github.com/huggingface/diffusers/issues/9165
+"""
+
+import os
+import torch
+import torch.nn as nn
+import bitsandbytes as bnb
+from transformers.quantizers.quantizers_utils import get_module_from_name
+from huggingface_hub import hf_hub_download
+from accelerate import init_empty_weights
+from accelerate.utils import set_module_tensor_to_device
+from diffusers import FluxTransformer2DModel, FluxPipeline
+from diffusers.loaders.single_file_utils import convert_flux_transformer_checkpoint_to_diffusers
+import safetensors.torch
+from modules import shared, devices
+
+
+def _replace_with_bnb_linear(
+    model,
+    method="nf4",
+    has_been_replaced=False,
+):
+    """
+    Private method that wraps the recursion for module replacement.
+
+    Returns the converted model and a boolean that indicates if the conversion has been successfull or not.
+    """
+    for name, module in model.named_children():
+        if isinstance(module, nn.Linear):
+            with init_empty_weights():
+                in_features = module.in_features
+                out_features = module.out_features
+
+                if method == "llm_int8":
+                    model._modules[name] = bnb.nn.Linear8bitLt( # pylint: disable=protected-access
+                        in_features,
+                        out_features,
+                        module.bias is not None,
+                        has_fp16_weights=False,
+                        threshold=6.0,
+                    )
+                    has_been_replaced = True
+                else:
+                    model._modules[name] = bnb.nn.Linear4bit( # pylint: disable=protected-access
+                        in_features,
+                        out_features,
+                        module.bias is not None,
+                        compute_dtype=torch.bfloat16,
+                        compress_statistics=False,
+                        quant_type="nf4",
+                    )
+                    has_been_replaced = True
+                # Store the module class in case we need to transpose the weight later
+                model._modules[name].source_cls = type(module) # pylint: disable=protected-access
+                # Force requires grad to False to avoid unexpected errors
+                model._modules[name].requires_grad_(False) # pylint: disable=protected-access
+
+        if len(list(module.children())) > 0:
+            _, has_been_replaced = _replace_with_bnb_linear(
+                module,
+                has_been_replaced=has_been_replaced,
+            )
+        # Remove the last key for recursion
+    return model, has_been_replaced
+
+
+def check_quantized_param(
+    model,
+    param_name: str,
+) -> bool:
+    module, tensor_name = get_module_from_name(model, param_name)
+    if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Params4bit): # pylint: disable=protected-access
+        # Add here check for loaded components' dtypes once serialization is implemented
+        return True
+    elif isinstance(module, bnb.nn.Linear4bit) and tensor_name == "bias":
+        # bias could be loaded by regular set_module_tensor_to_device() from accelerate,
+        # but it would wrongly use uninitialized weight there.
+        return True
+    else:
+        return False
+
+
+def create_quantized_param(
+    model,
+    param_value: "torch.Tensor",
+    param_name: str,
+    target_device: "torch.device",
+    state_dict=None,
+    unexpected_keys=None,
+    pre_quantized=False
+):
+    module, tensor_name = get_module_from_name(model, param_name)
+
+    if tensor_name not in module._parameters: # pylint: disable=protected-access
+        raise ValueError(f"{module} does not have a parameter or a buffer named {tensor_name}.")
+
+    old_value = getattr(module, tensor_name)
+
+    if tensor_name == "bias":
+        if param_value is None:
+            new_value = old_value.to(target_device)
+        else:
+            new_value = param_value.to(target_device)
+
+        new_value = torch.nn.Parameter(new_value, requires_grad=old_value.requires_grad)
+        module._parameters[tensor_name] = new_value # pylint: disable=protected-access
+        return
+
+    if not isinstance(module._parameters[tensor_name], bnb.nn.Params4bit): # pylint: disable=protected-access
+        raise ValueError("this function only loads `Linear4bit components`")
+    if (
+        old_value.device == torch.device("meta")
+        and target_device not in ["meta", torch.device("meta")]
+        and param_value is None
+    ):
+        raise ValueError(f"{tensor_name} is on the meta device, we need a `value` to put in on {target_device}.")
+
+    if pre_quantized:
+        if (param_name + ".quant_state.bitsandbytes__fp4" not in state_dict) and (
+                param_name + ".quant_state.bitsandbytes__nf4" not in state_dict
+            ):
+            raise ValueError(
+                f"Supplied state dict for {param_name} does not contain `bitsandbytes__*` and possibly other `quantized_stats` components."
+            )
+
+        quantized_stats = {}
+        for k, v in state_dict.items():
+            # `startswith` to counter for edge cases where `param_name`
+            # substring can be present in multiple places in the `state_dict`
+            if param_name + "." in k and k.startswith(param_name):
+                quantized_stats[k] = v
+                if unexpected_keys is not None and k in unexpected_keys:
+                    unexpected_keys.remove(k)
+
+        new_value = bnb.nn.Params4bit.from_prequantized(
+            data=param_value,
+            quantized_stats=quantized_stats,
+            requires_grad=False,
+            device=target_device,
+        )
+
+    else:
+        new_value = param_value.to("cpu")
+        kwargs = old_value.__dict__
+        new_value = bnb.nn.Params4bit(new_value, requires_grad=False, **kwargs).to(target_device)
+
+    module._parameters[tensor_name] = new_value # pylint: disable=protected-access
+
+
+def load_flux_nf4(checkpoint_info, diffusers_load_config):
+    if os.path.exists(checkpoint_info.path) and os.path.isfile(checkpoint_info.path):
+        ckpt_path = checkpoint_info.path
+    else:
+        ckpt_path = hf_hub_download(checkpoint_info.path, filename="diffusion_pytorch_model.safetensors", cache_dir=shared.opts.diffusers_dir)
+    original_state_dict = safetensors.torch.load_file(ckpt_path)
+
+    if 'sayakpaul/flux.1-dev-nf4' in checkpoint_info.path:
+        converted_state_dict = original_state_dict # already converted
+    else:
+        try:
+            converted_state_dict = convert_flux_transformer_checkpoint_to_diffusers(original_state_dict)
+        except Exception as e:
+            from modules import errors
+            errors.display(e, 'FLUX convert:')
+            raise
+
+    with init_empty_weights():
+        # config = FluxTransformer2DModel.load_config(checkpoint_info.path)
+        config = FluxTransformer2DModel.load_config("black-forest-labs/flux.1-dev", subfolder="transformer")
+        model = FluxTransformer2DModel.from_config(config).to(devices.dtype)
+        expected_state_dict_keys = list(model.state_dict().keys())
+
+    _replace_with_bnb_linear(model, "nf4")
+
+    for param_name, param in converted_state_dict.items():
+        if param_name not in expected_state_dict_keys:
+            continue
+        is_param_float8_e4m3fn = hasattr(torch, "float8_e4m3fn") and param.dtype == torch.float8_e4m3fn
+        if torch.is_floating_point(param) and not is_param_float8_e4m3fn:
+            param = param.to(devices.dtype)
+        if not check_quantized_param(model, param_name):
+            set_module_tensor_to_device(model, param_name, device=0, value=param)
+        else:
+            create_quantized_param(model, param, param_name, target_device=0, state_dict=original_state_dict, pre_quantized=True)
+
+    del original_state_dict
+    pipe = FluxPipeline.from_pretrained("black-forest-labs/flux.1-dev", transformer=model, cache_dir=shared.opts.diffusers_dir, **diffusers_load_config)
+    devices.torch_gc(force=True)
+    return pipe

modules/model_sd3.py

@@ -13,7 +13,7 @@ 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_major, diffusers_minor, diffusers_micro = int(diffusers.__version__.split('.')[0]), int(diffusers.__version__.split('.')[1]), int(diffusers.__version__.split('.')[2])
+        _diffusers_major, diffusers_minor, diffusers_micro = int(diffusers.__version__.split('.')[0]), int(diffusers.__version__.split('.')[1]), int(diffusers.__version__.split('.')[2]) # pylint: disable=use-maxsplit-arg
         fn_size = os.path.getsize(fn)
         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 = {

modules/processing_diffusers.py

@@ -314,7 +314,7 @@ def process_diffusers(p: processing.StableDiffusionProcessing):
                 shared.log.debug(f'Generated: frames={len(output.frames[0])}')
                 output.images = output.frames[0]
             if hasattr(shared.sd_model, "_unpack_latents") and hasattr(shared.sd_model, "vae_scale_factor"): # FLUX
-                output.images = shared.sd_model._unpack_latents(output.images, p.height, p.width, shared.sd_model.vae_scale_factor)
+                output.images = shared.sd_model._unpack_latents(output.images, p.height, p.width, shared.sd_model.vae_scale_factor) # pylint: disable=protected-access
             if torch.is_tensor(output.images) and len(output.images) > 0 and any(s >= 512 for s in output.images.shape):
                 results = output.images.float().cpu().numpy()
             elif hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0:

modules/sd_models.py

@@ -79,13 +79,15 @@ class CheckpointInfo:
             self.filename = filename
             self.sha256 = hashes.sha256_from_cache(self.filename, f"checkpoint/{relname}")
             self.type = ext
+            if 'nf4' in filename:
+                self.type = 'transformer'
         else: # maybe a diffuser
             if self.hash is None:
                 repo = [r for r in modelloader.diffuser_repos if self.filename == r['name']]
             else:
                 repo = [r for r in modelloader.diffuser_repos if self.hash == r['hash']]
             if len(repo) == 0:
-                self.name = relname
+                self.name = filename
                 self.filename = filename
                 self.sha256 = None
                 self.type = 'unknown'
@@ -707,7 +709,7 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model', offload=True):
         except Exception as e:
             shared.log.error(f'Error enabling fused projections: {e}')
     if shared.opts.diffusers_eval:
-        def eval_model(model, op=None, sd_model=None):
+        def eval_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
             if hasattr(model, "requires_grad_"):
                 model.requires_grad_(False)
                 model.eval()
@@ -782,7 +784,7 @@ def apply_balanced_offload(sd_model):
                 offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__))
                 module = dispatch_model(module, device_map=device_map, offload_dir=offload_dir)
                 module = add_hook_to_module(module, dispatch_from_cpu_hook(), append=True)
-                module._hf_hook.execution_device = torch.device(devices.device)
+                module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
             return args, kwargs
         def post_forward(self, module, output):
             return output
@@ -802,7 +804,7 @@ def apply_balanced_offload(sd_model):
                 module = module.to("cpu")
                 module.offload_dir = offload_dir
                 module = add_hook_to_module(module, dispatch_from_cpu_hook(), append=True)
-                module._hf_hook.execution_device = torch.device(devices.device)
+                module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access
                 devices.torch_gc()
 
     apply_balanced_offload_to_module(sd_model)
@@ -1029,7 +1031,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No
 
         shared.log.debug(f'Diffusers loading: path="{checkpoint_info.path}"')
         pipeline, model_type = detect_pipeline(checkpoint_info.path, op)
-        if os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface':
+        if os.path.isdir(checkpoint_info.path) or checkpoint_info.type == 'huggingface' or checkpoint_info.type == 'transformer':
             files = shared.walk_files(checkpoint_info.path, ['.safetensors', '.bin', '.ckpt'])
             if 'variant' not in diffusers_load_config and any('diffusion_pytorch_model.fp16' in f for f in files): # deal with diffusers lack of variant fallback when loading
                 diffusers_load_config['variant'] = 'fp16'

modules/sd_models_compile.py

@@ -96,7 +96,7 @@ def ipex_optimize(sd_model):
     try:
         t0 = time.time()
 
-        def ipex_optimize_model(model, op=None, sd_model=None):
+        def ipex_optimize_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
             import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
             model.eval()
             model.training = False
@@ -133,7 +133,7 @@ def nncf_send_to_device(model):
 
 def nncf_compress_model(model, op=None, sd_model=None):
     import nncf
-    global quant_last_model_name, quant_last_model_device
+    global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
     model.eval()
     backup_embeddings = None
     if hasattr(model, "get_input_embeddings"):
@@ -164,7 +164,7 @@ def nncf_compress_weights(sd_model):
     try:
         t0 = time.time()
         shared.log.info(f"NNCF Compress Weights: {shared.opts.nncf_compress_weights}")
-        global quant_last_model_name, quant_last_model_device
+        global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
         from installer import install
         install('nncf==2.7.0', quiet=True)
 
@@ -186,8 +186,8 @@ def nncf_compress_weights(sd_model):
     return sd_model
 
 def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activations=None):
-    from optimum import quanto
-    global quant_last_model_name, quant_last_model_device
+    from optimum import quanto # pylint: disable=no-name-in-module
+    global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
     if sd_model is not None and "Flux" in sd_model.__class__.__name__: # LayerNorm is not supported
         exclude_list = ["transformer_blocks.*.norm1.norm", "transformer_blocks.*.norm2", "transformer_blocks.*.norm1_context.norm", "transformer_blocks.*.norm2_context", "single_transformer_blocks.*.norm.norm", "norm_out.norm"]
     else:
@@ -228,14 +228,14 @@ def optimum_quanto_model(model, op=None, sd_model=None, weights=None, activation
 def optimum_quanto_weights(sd_model):
     try:
         if shared.opts.diffusers_offload_mode in {"balanced", "sequential"}:
-            shared.log.warn(f"Optimum Quanto Weights is incompatible with {shared.opts.diffusers_offload_mode} offload!")
+            shared.log.warning(f"Optimum Quanto Weights is incompatible with {shared.opts.diffusers_offload_mode} offload!")
             return sd_model
         t0 = time.time()
         shared.log.info(f"Optimum Quanto Weights: {shared.opts.optimum_quanto_weights}")
-        global quant_last_model_name, quant_last_model_device
+        global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement
         from installer import install
         install('optimum-quanto', quiet=True)
-        from optimum import quanto
+        from optimum import quanto # pylint: disable=no-name-in-module
         quanto.tensor.qbits.QBitsTensor.create = lambda *args, **kwargs: quanto.tensor.qbits.QBitsTensor(*args, **kwargs)
 
         sd_model = apply_compile_to_model(sd_model, optimum_quanto_model, shared.opts.optimum_quanto_weights, op="optimum-quanto")
@@ -255,7 +255,7 @@ def optimum_quanto_weights(sd_model):
             activations = None
 
         if activations is not None:
-            def optimum_quanto_freeze(model, op=None, sd_model=None):
+            def optimum_quanto_freeze(model, op=None, sd_model=None): # pylint: disable=unused-argument
                 quanto.freeze(model)
                 return model
             if shared.opts.diffusers_offload_mode == "model":
@@ -382,7 +382,7 @@ def compile_torch(sd_model):
         torch._dynamo.reset() # pylint: disable=protected-access
         shared.log.debug(f"Model compile available backends: {torch._dynamo.list_backends()}") # pylint: disable=protected-access
 
-        def torch_compile_model(model, op=None, sd_model=None):
+        def torch_compile_model(model, op=None, sd_model=None): # pylint: disable=unused-argument
             if model.device.type != "meta":
                 return_device = model.device
                 model = torch.compile(model.to(devices.device),

modules/sd_samplers.py

@@ -69,6 +69,10 @@ def create_sampler(name, model):
         return sampler
     elif shared.native:
         sampler = config.constructor(model)
+        if shared.sd_model_type == 'FluxPipeline':
+            if 'base_image_seq_len' not in sampler.sampler.config or 'max_image_seq_len' not in sampler.sampler.config or 'base_shift' not in sampler.sampler.config or 'max_shift' not in sampler.sampler.config:
+                shared.log.warning('FLUX sampler: attempting to use a non compatible scheduler')
+                return None
         if not hasattr(model, 'scheduler_config'):
             model.scheduler_config = sampler.sampler.config.copy()
         model.scheduler = sampler.sampler

modules/shared.py

@@ -3,12 +3,12 @@ import os
 import sys
 import time
 import json
-import psutil
 import threading
 import contextlib
 from types import SimpleNamespace
 from urllib.parse import urlparse
 from enum import Enum
+import psutil
 import requests
 import gradio as gr
 import fasteners

modules/zluda_installer.py

@@ -24,14 +24,7 @@ def install(zluda_path: os.PathLike) -> None:
     if os.path.exists(zluda_path):
         return
 
-    default_hash = None
-    if rocm.version == "6.1":
-        default_hash = '2f2e38a8adebb456ad75390e60871f2c8ba18fa7'
-    elif rocm.version == "5.7":
-        default_hash = '11cc5844514f93161e0e74387f04e2c537705a82'
-    else:
-        raise RuntimeError(f'Unsupported HIP SDK version: {rocm.version}')
-    urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{os.environ.get("ZLUDA_HASH", default_hash)}/ZLUDA-windows-amd64.zip', '_zluda')
+    urllib.request.urlretrieve(f'https://github.com/lshqqytiger/ZLUDA/releases/download/rel.{os.environ.get("ZLUDA_HASH", "1c238a959f2aafdb9900f6801b61d9c0318040a2")}/ZLUDA-windows-rocm{rocm.version[0]}-amd64.zip', '_zluda')
     with zipfile.ZipFile('_zluda', 'r') as archive:
         infos = archive.infolist()
         for info in infos: