From 3868a9184baf82ab70bce15f69a17258084edfe0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Enes=20Sad=C4=B1k=20=C3=96zbek?= <1809172+Trojaner@users.noreply.github.com> Date: Sat, 7 Jun 2025 00:36:08 +0300 Subject: [PATCH 01/25] Fix incorrectly reported lycoris load error --- modules/lora/lora_load.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/modules/lora/lora_load.py b/modules/lora/lora_load.py index f02e4a274..dbeb57630 100644 --- a/modules/lora/lora_load.py +++ b/modules/lora/lora_load.py @@ -148,8 +148,9 @@ def load_safetensors(name, network_on_disk) -> Union[network.Network, None]: if net_module is not None: network_types.append(nettype.__class__.__name__) break - module_errors += 1 if net_module is None: + module_errors += 1 + if l.debug: shared.log.error(f'LoRA unhandled: name={name} key={key} weights={weights.w.keys()}') else: From 86cd272b961af0f317f4faf4b7b6e5046e4c3b76 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 18 Jun 2025 16:24:42 +0300 Subject: [PATCH 02/25] SDNQ fix Dora --- CHANGELOG.md | 9 +++++++++ modules/lora/lora_apply.py | 5 ++++- modules/lora/network.py | 5 ++++- modules/model_quant.py | 2 +- modules/sdnq/__init__.py | 2 ++ modules/sdnq/dequantizer.py | 8 ++++++++ 6 files changed, 28 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ed2ebb762..f11448519 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,14 @@ # Change Log for SD.Next +## Update for 2025-06-18 + +- **SDNQ Quantization** + - Fix Qwen 2.5 with int8 matmul + - Fix Dora loading + +- **Fixes** + - IPEX with DPM2++ FlowMatch samplers + ## Update for 2025-06-16 - **Feature** diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py index b971541cd..2fcea174c 100644 --- a/modules/lora/lora_apply.py +++ b/modules/lora/lora_apply.py @@ -79,7 +79,9 @@ def network_calc_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn. continue try: t0 = time.time() - if hasattr(self, "sdnq_dequantizer"): + if hasattr(self, "sdnq_dequantizer_backup"): + weight = self.sdnq_dequantizer_backup.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer_backup.use_quantized_matmul) + elif hasattr(self, "sdnq_dequantizer"): weight = self.sdnq_dequantizer.to(devices.device)(self.weight.to(devices.device), skip_quantized_matmul=self.sdnq_dequantizer.use_quantized_matmul) else: weight = self.weight.to(devices.device) # must perform calc on gpu due to performance @@ -228,6 +230,7 @@ def network_apply_weights(self: Union[torch.nn.Conv2d, torch.nn.Linear, torch.nn self.weight = torch.nn.Parameter(weights_backup.to(device), requires_grad=False) if hasattr(self, "sdnq_dequantizer_backup"): self.sdnq_dequantizer = self.sdnq_dequantizer_backup.to(device) + del self.sdnq_dequantizer_backup if bias_backup is not None: self.bias = None diff --git a/modules/lora/network.py b/modules/lora/network.py index f6d93009c..44e5a64b9 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -124,7 +124,10 @@ class NetworkModule: self.sd_key = weights.sd_key self.sd_module = weights.sd_module if hasattr(self.sd_module, 'weight'): - self.shape = self.sd_module.weight.shape + if hasattr(self.sd_module, "sdnq_dequantizer"): + self.shape = self.sd_module.sdnq_dequantizer.original_shape + else: + self.shape = self.sd_module.weight.shape self.dim = None self.bias = weights.w.get("bias") self.alpha = weights.w["alpha"].item() if "alpha" in weights.w else None diff --git a/modules/model_quant.py b/modules/model_quant.py index 3c57ac67c..75632c020 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -402,7 +402,7 @@ def sdnq_quantize_weights(sd_model): try: t0 = time.time() from modules import shared, devices, sd_models - log.info(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights}") + log.info(f"Quantization: type=SDNQ dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} modules={shared.opts.sdnq_quantize_weights}") global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq") diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index 998e0a467..cd9fc7759 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -21,6 +21,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz is_conv_transpose_type = False is_linear_type = False result_shape = None + original_shape = layer.weight.shape if torch_dtype is None: torch_dtype = devices.dtype @@ -139,6 +140,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz quantized_weight_shape=layer.weight.shape, result_dtype=torch_dtype, result_shape=result_shape, + original_shape=original_shape, weights_dtype=weights_dtype, use_quantized_matmul=use_quantized_matmul, ) diff --git a/modules/sdnq/dequantizer.py b/modules/sdnq/dequantizer.py index 114692f75..45cedfa79 100644 --- a/modules/sdnq/dequantizer.py +++ b/modules/sdnq/dequantizer.py @@ -46,11 +46,13 @@ class AsymmetricWeightsDequantizer(torch.nn.Module): zero_point: torch.FloatTensor, result_dtype: torch.dtype, result_shape: torch.Size, + original_shape: torch.Size, weights_dtype: str, **kwargs, # pylint: disable=unused-argument ): super().__init__() self.weights_dtype = weights_dtype + self.original_shape = original_shape self.use_quantized_matmul = False self.result_dtype = result_dtype self.result_shape = result_shape @@ -70,12 +72,14 @@ class SymmetricWeightsDequantizer(torch.nn.Module): scale: torch.FloatTensor, result_dtype: torch.dtype, result_shape: torch.Size, + original_shape: torch.Size, weights_dtype: str, use_quantized_matmul: bool = False, **kwargs, # pylint: disable=unused-argument ): super().__init__() self.weights_dtype = weights_dtype + self.original_shape = original_shape self.use_quantized_matmul = use_quantized_matmul self.result_dtype = result_dtype self.result_shape = result_shape @@ -96,12 +100,14 @@ class PackedINTAsymmetricWeightsDequantizer(torch.nn.Module): quantized_weight_shape: torch.Size, result_dtype: torch.dtype, result_shape: torch.Size, + original_shape: torch.Size, weights_dtype: str, **kwargs, # pylint: disable=unused-argument ): super().__init__() self.weights_dtype = weights_dtype self.use_quantized_matmul = False + self.original_shape = original_shape self.quantized_weight_shape = quantized_weight_shape self.result_dtype = result_dtype self.result_shape = result_shape @@ -122,12 +128,14 @@ class PackedINTSymmetricWeightsDequantizer(torch.nn.Module): quantized_weight_shape: torch.Size, result_dtype: torch.dtype, result_shape: torch.Size, + original_shape: torch.Size, weights_dtype: str, use_quantized_matmul: bool = False, **kwargs, # pylint: disable=unused-argument ): super().__init__() self.weights_dtype = weights_dtype + self.original_shape = original_shape self.use_quantized_matmul = use_quantized_matmul self.quantized_weight_shape = quantized_weight_shape self.result_dtype = result_dtype From 2c4850cc2ba771364d8ffc91d340858e2c4bae4e Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 18 Jun 2025 18:03:46 +0300 Subject: [PATCH 03/25] Log more info with SDNQ --- modules/model_quant.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index 75632c020..ac8186cce 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -142,7 +142,7 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo quantization_device=quantization_device, return_device=return_device, ) - log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype}') + log.debug(f'Quantization: module="{module}" type=sdnq dtype={weights_dtype} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} quantization_device={quantization_device} return_device={return_device}') if kwargs is None: return sdnq_config else: @@ -402,7 +402,7 @@ def sdnq_quantize_weights(sd_model): try: t0 = time.time() from modules import shared, devices, sd_models - log.info(f"Quantization: type=SDNQ dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} modules={shared.opts.sdnq_quantize_weights}") + log.info(f"Quantization: type=SDNQ dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} modules={shared.opts.sdnq_quantize_weights}") global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq") From ff5d7977a9c6ffc8e9bc8a229c9833547c40bdd8 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 18 Jun 2025 19:35:57 +0300 Subject: [PATCH 04/25] Fix Controlnet pipeline check on set atten --- modules/sd_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/sd_models.py b/modules/sd_models.py index d8aad13a3..88644b17e 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -920,7 +920,7 @@ def set_diffusers_attention(pipe, quiet:bool=False): # if hasattr(pipe, 'pipe'): # set_diffusers_attention(pipe.pipe) - if 'ControlNet' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")): + if 'Control' in pipe.__class__.__name__ or 'Adapter' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")): if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}: shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}") else: From 7fc7797a1dd536d03516f07b7632306ba1f54380 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 19 Jun 2025 11:36:41 +0300 Subject: [PATCH 05/25] SDNQ fix group size calc on odd shapes --- modules/sdnq/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index cd9fc7759..b58721004 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -77,13 +77,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz num_of_groups = 1 else: num_of_groups = channel_size // group_size - while channel_size % group_size != 0: # find something divisible + while num_of_groups * group_size != channel_size: # find something divisible num_of_groups -= 1 if num_of_groups <= 1: group_size = channel_size num_of_groups = 1 break - group_size = channel_size / num_of_groups + group_size = channel_size // num_of_groups group_size = int(group_size) num_of_groups = int(num_of_groups) From 4491d26a184818fad3cf7cbbd1d7d1b801f10184 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 19 Jun 2025 22:55:28 +0300 Subject: [PATCH 06/25] IPEX fix fp64 hijack --- modules/intel/ipex/hijacks.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/intel/ipex/hijacks.py b/modules/intel/ipex/hijacks.py index 2a3b9e06a..e2a04a662 100644 --- a/modules/intel/ipex/hijacks.py +++ b/modules/intel/ipex/hijacks.py @@ -259,7 +259,7 @@ def Tensor_to(self, device=None, *args, **kwargs): return self.original_Tensor_to(return_xpu(device), *args, **kwargs) else: if not device_supports_fp64: - if kwargs.get("dtype", None) == torch.float64 and torch.device(device).type == "xpu": + if kwargs.get("dtype", None) == torch.float64 and ((device is None and self.device.type == "xpu") or (device is not None and torch.device(device).type == "xpu")): kwargs["dtype"] = torch.float32 elif device == torch.float64 and self.device.type == "xpu": device = torch.float32 From 87e6d3f4fc9da29b2ee738951c1ee3c892436112 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Fri, 20 Jun 2025 20:41:11 +0300 Subject: [PATCH 07/25] SDNQ add modules_to_not_convert and don't quant _keep_in_fp32_modules layers in post mode --- modules/model_quant.py | 15 +++++---------- modules/sdnq/__init__.py | 25 ++++--------------------- 2 files changed, 9 insertions(+), 31 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index ac8186cce..b3db3f311 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -328,16 +328,6 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): from modules.sdnq import apply_sdnq_to_module model.eval() - - if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"}: - import torch - from modules.sdnq import SDNQ_T5DenseGatedActDense # T5DenseGatedActDense uses fp32 - for i in range(len(model.encoder.block)): - model.encoder.block[i].layer[1].DenseReluDense = SDNQ_T5DenseGatedActDense( - model.encoder.block[i].layer[1].DenseReluDense, - dtype=torch.float32 if devices.dtype != torch.bfloat16 else torch.bfloat16 - ) - backup_embeddings = None if hasattr(model, "get_input_embeddings"): backup_embeddings = copy.deepcopy(model.get_input_embeddings()) @@ -357,6 +347,10 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): quantization_device = None return_device = None + modules_to_not_convert = getattr(model, "_keep_in_fp32_modules", []) + if modules_to_not_convert is None: + modules_to_not_convert = [] + model = apply_sdnq_to_module( model, weights_dtype=weights_dtype, @@ -369,6 +363,7 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): quantization_device=quantization_device, return_device=return_device, param_name=op, + modules_to_not_convert=modules_to_not_convert, ) model.quantization_method = 'SDNQ' diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index b58721004..ceb19366c 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -153,11 +153,13 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz return layer -def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None): # pylint: disable=unused-argument +def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_size=0, quant_conv=False, use_quantized_matmul=False, use_quantized_matmul_conv=False, dequantize_fp32=False, quantization_device=None, return_device=None, param_name=None, modules_to_not_convert: List[str] = []): # pylint: disable=unused-argument has_children = list(model.children()) if not has_children: return model for module_param_name, module in model.named_children(): + if module_param_name in modules_to_not_convert: + continue if hasattr(module, "weight") and module.weight is not None: module = sdnq_quantize_layer( module, @@ -184,6 +186,7 @@ def apply_sdnq_to_module(model, weights_dtype="int8", torch_dtype=None, group_si quantization_device=quantization_device, return_device=return_device, param_name=module_param_name, + modules_to_not_convert=modules_to_not_convert, ) return model @@ -440,23 +443,3 @@ class SDNQConfig(QuantizationConfigMixin): raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] - - -class SDNQ_T5DenseGatedActDense(torch.nn.Module): # forward can't find what self is without creating a class - def __init__(self, T5DenseGatedActDense, dtype): - super().__init__() - self.wi_0 = T5DenseGatedActDense.wi_0 - self.wi_1 = T5DenseGatedActDense.wi_1 - self.wo = T5DenseGatedActDense.wo - self.dropout = T5DenseGatedActDense.dropout - self.act = T5DenseGatedActDense.act - self.torch_dtype = dtype - - def forward(self, hidden_states): - hidden_gelu = self.act(self.wi_0(hidden_states)) - hidden_linear = self.wi_1(hidden_states) - hidden_states = hidden_gelu * hidden_linear - hidden_states = self.dropout(hidden_states) - hidden_states = hidden_states.to(self.torch_dtype) # this line needs to be forced - hidden_states = self.wo(hidden_states) - return hidden_states From bbf986d3a581fd2b357341bba53d8d155a4b1c7a Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 23 Jun 2025 17:37:20 +0300 Subject: [PATCH 08/25] Fix LTXVideo --- CHANGELOG.md | 4 +++- modules/sdnq/__init__.py | 2 +- modules/video_models/video_load.py | 4 ++++ 3 files changed, 8 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f11448519..a9938fed8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,13 +1,15 @@ # Change Log for SD.Next -## Update for 2025-06-18 +## Update for 2025-06-123 - **SDNQ Quantization** - Fix Qwen 2.5 with int8 matmul - Fix Dora loading + - Remove per layer GC - **Fixes** - IPEX with DPM2++ FlowMatch samplers + - LTXVideo default scheduler ## Update for 2025-06-16 diff --git a/modules/sdnq/__init__.py b/modules/sdnq/__init__.py index ceb19366c..c76d1e9f9 100644 --- a/modules/sdnq/__init__.py +++ b/modules/sdnq/__init__.py @@ -149,7 +149,7 @@ def sdnq_quantize_layer(layer, weights_dtype="int8", torch_dtype=None, group_siz layer.forward = get_forward_func(layer_class_name, use_quantized_matmul, dtype_dict[weights_dtype]["is_integer"], use_tensorwise_fp8_matmul) layer.forward = layer.forward.__get__(layer, layer.__class__) - devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}") + #devices.torch_gc(force=False, reason=f"SDNQ param_name: {param_name}") return layer diff --git a/modules/video_models/video_load.py b/modules/video_models/video_load.py index 33813df82..d607b8fc3 100644 --- a/modules/video_models/video_load.py +++ b/modules/video_models/video_load.py @@ -1,4 +1,5 @@ import os +import copy import time from modules import shared, errors, sd_models, sd_checkpoint, model_quant, devices, sd_hijack_te from modules.video_models import models_def, video_utils, video_vae, video_overrides, video_cache @@ -70,6 +71,9 @@ def load_model(selected: models_def.Model): errors.display(e, 'video') t1 = time.time() + if shared.sd_model.__class__.__name__.startswith("LTX"): + shared.sd_model.scheduler.config.use_dynamic_shifting = False + shared.sd_model.default_scheduler = copy.deepcopy(shared.sd_model.scheduler) if hasattr(shared.sd_model, "scheduler") else None shared.sd_model.sd_checkpoint_info = sd_checkpoint.CheckpointInfo(selected.repo) shared.sd_model.sd_model_hash = None sd_models.set_diffuser_options(shared.sd_model) From 8ef348d9edb9a972cad74692aac0af22647a8df5 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Mon, 23 Jun 2025 17:39:39 +0300 Subject: [PATCH 09/25] Update changelog --- CHANGELOG.md | 4 +++- wiki | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a9938fed8..0c6094039 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,14 +1,16 @@ # Change Log for SD.Next -## Update for 2025-06-123 +## Update for 2025-06-23 - **SDNQ Quantization** + - Add modules_to_not_convert support for post mode - Fix Qwen 2.5 with int8 matmul - Fix Dora loading - Remove per layer GC - **Fixes** - IPEX with DPM2++ FlowMatch samplers + - Invalid attention processor with ControlNet - LTXVideo default scheduler ## Update for 2025-06-16 diff --git a/wiki b/wiki index 2ca67daca..5e97702f2 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 2ca67dacaccb7ef4c0595b19407fc5a93167008b +Subproject commit 5e97702f219b879c035057204303ae649e1edcf7 From eaa130463cd0a2a20f6165153566d5a26f8977b9 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 24 Jun 2025 13:09:11 +0300 Subject: [PATCH 10/25] Use Diffusers with OmniGen --- CHANGELOG.md | 8 +- html/reference.json | 2 +- modules/model_omnigen.py | 60 +++-- modules/omnigen/__init__.py | 4 - modules/omnigen/model.py | 390 --------------------------------- modules/omnigen/pipeline.py | 219 ------------------ modules/omnigen/processor.py | 312 -------------------------- modules/omnigen/scheduler.py | 55 ----- modules/omnigen/transformer.py | 164 -------------- modules/omnigen/utils.py | 105 --------- modules/processing_args.py | 2 +- modules/sd_offload.py | 2 +- modules/sd_vae_remote.py | 15 +- modules/sd_vae_taesd.py | 4 +- modules/shared_items.py | 2 +- 15 files changed, 67 insertions(+), 1277 deletions(-) delete mode 100644 modules/omnigen/__init__.py delete mode 100644 modules/omnigen/model.py delete mode 100644 modules/omnigen/pipeline.py delete mode 100644 modules/omnigen/processor.py delete mode 100644 modules/omnigen/scheduler.py delete mode 100644 modules/omnigen/transformer.py delete mode 100644 modules/omnigen/utils.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c6094039..5faadfa04 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Change Log for SD.Next -## Update for 2025-06-23 +## Update for 2025-06-24 + +- **Changes** + - Use Diffusers version of OmniGen + - Support Remote VAE with Omnigen, Lumina 2 and PixArt - **SDNQ Quantization** - Add modules_to_not_convert support for post mode @@ -12,6 +16,8 @@ - IPEX with DPM2++ FlowMatch samplers - Invalid attention processor with ControlNet - LTXVideo default scheduler + - Balanced offload with OmniGen + - Quantization with OmniGen ## Update for 2025-06-16 diff --git a/html/reference.json b/html/reference.json index 2441112e6..fe89f47ac 100644 --- a/html/reference.json +++ b/html/reference.json @@ -239,7 +239,7 @@ }, "VectorSpaceLab OmniGen v1": { - "path": "Shitao/OmniGen-v1", + "path": "Shitao/OmniGen-v1-diffusers", "desc": "OmniGen is a unified image generation model that can generate a wide range of images from multi-modal prompts. It is designed to be simple, flexible and easy to use.", "preview": "Shitao--OmniGen-v1.jpg", "skip": true diff --git a/modules/model_omnigen.py b/modules/model_omnigen.py index b7eb4684e..0df4948a6 100644 --- a/modules/model_omnigen.py +++ b/modules/model_omnigen.py @@ -1,25 +1,47 @@ -def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument - from modules import shared, devices, sd_models, shared_items - repo_id = sd_models.path_to_repo(checkpoint_info.name) +import os +import diffusers +from modules import errors, shared, devices, sd_models, model_quant - # load - from modules.omnigen import OmniGenPipeline - shared_items.pipelines['OmniGen'] = OmniGenPipeline - pipe = OmniGenPipeline.from_pretrained( - model_name=repo_id, - vae_path='madebyollin/sdxl-vae-fp16-fix', +debug = shared.log.trace if os.environ.get('SD_LOAD_DEBUG', None) is not None else lambda *args, **kwargs: None + + +def load_omnigen(checkpoint_info, diffusers_load_config={}): # pylint: disable=unused-argument + repo_id = sd_models.path_to_repo(checkpoint_info.name) + vae = None + + if shared.opts.sd_vae != 'Default' and shared.opts.sd_vae != 'Automatic': + try: + debug(f'Load model: type=OmniGen vae="{shared.opts.sd_vae}"') + from modules import sd_vae + # vae = sd_vae.load_vae_diffusers(None, sd_vae.vae_dict[shared.opts.sd_vae], 'override') + vae_file = sd_vae.vae_dict[shared.opts.sd_vae] + if os.path.exists(vae_file): + vae_config = os.path.join('configs', 'sdxl', 'vae', 'config.json') + vae = diffusers.AutoencoderKL.from_single_file(vae_file, config=vae_config, **diffusers_load_config) + except Exception as e: + shared.log.error(f"Load model: type=OmniGen failed to load VAE: {e}") + shared.opts.sd_vae = 'Default' + if debug: + errors.display(e, 'OmniGen VAE:') + + load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, module='Transformer') + transformer = diffusers.OmniGenTransformer2DModel.from_pretrained( + repo_id, + subfolder="transformer", cache_dir=shared.opts.diffusers_dir, + **load_config, + **quant_config, + ) + + load_config, quant_config = model_quant.get_dit_args(diffusers_load_config, allow_quant=False) + if vae is not None: + load_config['vae'] = vae + pipe = diffusers.OmniGenPipeline.from_pretrained( + repo_id, + transformer=transformer, + cache_dir=shared.opts.diffusers_dir, + **load_config, ) - # init - pipe.device = devices.device - pipe.dtype = devices.dtype - pipe.model.device = devices.device - pipe.separate_cfg_infer = True - pipe.use_kv_cache = False - pipe.model.to(device=devices.device, dtype=devices.dtype) - if shared.opts.diffusers_eval: - pipe.model.eval() - pipe.vae.to(devices.device, dtype=devices.dtype) devices.torch_gc(force=True) return pipe diff --git a/modules/omnigen/__init__.py b/modules/omnigen/__init__.py deleted file mode 100644 index 40315a6f3..000000000 --- a/modules/omnigen/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from .model import OmniGen -from .processor import OmniGenProcessor -from .scheduler import OmniGenScheduler -from .pipeline import OmniGenPipeline diff --git a/modules/omnigen/model.py b/modules/omnigen/model.py deleted file mode 100644 index 17d696b53..000000000 --- a/modules/omnigen/model.py +++ /dev/null @@ -1,390 +0,0 @@ -# The code is revised from DiT -import os -import math -import torch -import torch.nn as nn -import numpy as np -from safetensors.torch import load_file -from diffusers.loaders import PeftAdapterMixin -from huggingface_hub import snapshot_download -from .transformer import Phi3Config, Phi3Transformer - - -def modulate(x, shift, scale): - return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1) - - -class TimestepEmbedder(nn.Module): - """ - Embeds scalar timesteps into vector representations. - """ - def __init__(self, hidden_size, frequency_embedding_size=256): - super().__init__() - self.mlp = nn.Sequential( - nn.Linear(frequency_embedding_size, hidden_size, bias=True), - nn.SiLU(), - nn.Linear(hidden_size, hidden_size, bias=True), - ) - self.frequency_embedding_size = frequency_embedding_size - - @staticmethod - def timestep_embedding(t, dim, max_period=10000): - """ - Create sinusoidal timestep embeddings. - :param t: a 1-D Tensor of N indices, one per batch element. - These may be fractional. - :param dim: the dimension of the output. - :param max_period: controls the minimum frequency of the embeddings. - :return: an (N, D) Tensor of positional embeddings. - """ - # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py - half = dim // 2 - freqs = torch.exp( - -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half - ).to(device=t.device) - args = t[:, None].float() * freqs[None] - embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) - if dim % 2: - embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) - return embedding - - def forward(self, t, dtype=torch.float32): - t_freq = self.timestep_embedding(t, self.frequency_embedding_size).to(dtype) - t_emb = self.mlp(t_freq) - return t_emb - - -class FinalLayer(nn.Module): - """ - The final layer of DiT. - """ - def __init__(self, hidden_size, patch_size, out_channels): - super().__init__() - self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) - self.linear = nn.Linear(hidden_size, patch_size * patch_size * out_channels, bias=True) - self.adaLN_modulation = nn.Sequential( - nn.SiLU(), - nn.Linear(hidden_size, 2 * hidden_size, bias=True) - ) - - def forward(self, x, c): - shift, scale = self.adaLN_modulation(c).chunk(2, dim=1) - x = modulate(self.norm_final(x), shift, scale) - x = self.linear(x) - return x - - -def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False, extra_tokens=0, interpolation_scale=1.0, base_size=1): - """ - grid_size: int of the grid height and width return: pos_embed: [grid_size*grid_size, embed_dim] or - [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token) - """ - if isinstance(grid_size, int): - grid_size = (grid_size, grid_size) - - grid_h = np.arange(grid_size[0], dtype=np.float32) / (grid_size[0] / base_size) / interpolation_scale - grid_w = np.arange(grid_size[1], dtype=np.float32) / (grid_size[1] / base_size) / interpolation_scale - grid = np.meshgrid(grid_w, grid_h) # here w goes first - grid = np.stack(grid, axis=0) - - grid = grid.reshape([2, 1, grid_size[1], grid_size[0]]) - pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid) - if cls_token and extra_tokens > 0: - pos_embed = np.concatenate([np.zeros([extra_tokens, embed_dim]), pos_embed], axis=0) - return pos_embed - - -def get_2d_sincos_pos_embed_from_grid(embed_dim, grid): - assert embed_dim % 2 == 0 - - # use half of dimensions to encode grid_h - emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2) - emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2) - - emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D) - return emb - - -def get_1d_sincos_pos_embed_from_grid(embed_dim, pos): - """ - embed_dim: output dimension for each position - pos: a list of positions to be encoded: size (M,) - out: (M, D) - """ - assert embed_dim % 2 == 0 - omega = np.arange(embed_dim // 2, dtype=np.float64) - omega /= embed_dim / 2. - omega = 1. / 10000**omega # (D/2,) - - pos = pos.reshape(-1) # (M,) - out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product - - emb_sin = np.sin(out) # (M, D/2) - emb_cos = np.cos(out) # (M, D/2) - - emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D) - return emb - - -class PatchEmbedMR(nn.Module): - """ 2D Image to Patch Embedding - """ - def __init__( - self, - patch_size: int = 2, - in_chans: int = 4, - embed_dim: int = 768, - bias: bool = True, - ): - super().__init__() - self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias) - - def forward(self, x): - x = self.proj(x) - x = x.flatten(2).transpose(1, 2) # NCHW -> NLC - return x - - -class OmniGen(nn.Module, PeftAdapterMixin): - """ - Diffusion model with a Transformer backbone. - """ - def __init__( - self, - transformer_config: Phi3Config, - patch_size=2, - in_channels=4, - pe_interpolation: float = 1.0, - pos_embed_max_size: int = 192, - ): - super().__init__() - self.in_channels = in_channels - self.out_channels = in_channels - self.patch_size = patch_size - self.pos_embed_max_size = pos_embed_max_size - hidden_size = transformer_config.hidden_size - self.x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True) - self.input_x_embedder = PatchEmbedMR(patch_size, in_channels, hidden_size, bias=True) - self.time_token = TimestepEmbedder(hidden_size) - self.t_embedder = TimestepEmbedder(hidden_size) - self.pe_interpolation = pe_interpolation - pos_embed = get_2d_sincos_pos_embed(hidden_size, pos_embed_max_size, interpolation_scale=self.pe_interpolation, base_size=64) - self.register_buffer("pos_embed", torch.from_numpy(pos_embed).float().unsqueeze(0), persistent=True) - self.final_layer = FinalLayer(hidden_size, patch_size, self.out_channels) - self.initialize_weights() - self.llm = Phi3Transformer(config=transformer_config) - self.llm.config.use_cache = False - - @classmethod - def from_pretrained(cls, model_name: str, cache_dir: str=None): - if not os.path.exists(os.path.join(model_name, 'model.pt')) and not os.path.exists(os.path.join(model_name, 'model.safetensors')): - cache_dir = cache_dir or os.getenv('HF_HUB_CACHE') - model_name = snapshot_download(repo_id=model_name, - cache_dir=cache_dir, - ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5']) - config = Phi3Config.from_pretrained(model_name) - model = cls(config) - if os.path.exists(os.path.join(model_name, 'model.pt')): - state_dict = torch.load(os.path.join(model_name, 'model.pt'), map_location='cpu') - elif os.path.exists(os.path.join(model_name, 'model.safetensors')): - state_dict = load_file(os.path.join(model_name, 'model.safetensors')) - else: - raise ValueError(f"OmniGen: Could not find model file in {model_name}") - model.load_state_dict(state_dict) - return model - - def initialize_weights(self): - assert not hasattr(self, "llama") - - # Initialize transformer layers: - def _basic_init(module): - if isinstance(module, nn.Linear): - torch.nn.init.xavier_uniform_(module.weight) - if module.bias is not None: - nn.init.constant_(module.bias, 0) - self.apply(_basic_init) - - # Initialize patch_embed like nn.Linear (instead of nn.Conv2d): - w = self.x_embedder.proj.weight.data - nn.init.xavier_uniform_(w.view([w.shape[0], -1])) - nn.init.constant_(self.x_embedder.proj.bias, 0) - - w = self.input_x_embedder.proj.weight.data - nn.init.xavier_uniform_(w.view([w.shape[0], -1])) - nn.init.constant_(self.x_embedder.proj.bias, 0) - - - # Initialize timestep embedding MLP: - nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02) - nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02) - nn.init.normal_(self.time_token.mlp[0].weight, std=0.02) - nn.init.normal_(self.time_token.mlp[2].weight, std=0.02) - - # Zero-out output layers: - nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0) - nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0) - nn.init.constant_(self.final_layer.linear.weight, 0) - nn.init.constant_(self.final_layer.linear.bias, 0) - - def unpatchify(self, x, h, w): - """ - x: (N, T, patch_size**2 * C) - imgs: (N, H, W, C) - """ - c = self.out_channels - - x = x.reshape(shape=(x.shape[0], h//self.patch_size, w//self.patch_size, self.patch_size, self.patch_size, c)) - x = torch.einsum('nhwpqc->nchpwq', x) - imgs = x.reshape(shape=(x.shape[0], c, h, w)) - return imgs - - - def cropped_pos_embed(self, height, width): - """Crops positional embeddings for SD3 compatibility.""" - if self.pos_embed_max_size is None: - raise ValueError("`pos_embed_max_size` must be set for cropping.") - - height = height // self.patch_size - width = width // self.patch_size - if height > self.pos_embed_max_size: - raise ValueError( - f"Height ({height}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - if width > self.pos_embed_max_size: - raise ValueError( - f"Width ({width}) cannot be greater than `pos_embed_max_size`: {self.pos_embed_max_size}." - ) - - top = (self.pos_embed_max_size - height) // 2 - left = (self.pos_embed_max_size - width) // 2 - spatial_pos_embed = self.pos_embed.reshape(1, self.pos_embed_max_size, self.pos_embed_max_size, -1) - spatial_pos_embed = spatial_pos_embed[:, top : top + height, left : left + width, :] - spatial_pos_embed = spatial_pos_embed.reshape(1, -1, spatial_pos_embed.shape[-1]) - return spatial_pos_embed - - - def patch_multiple_resolutions(self, latents, padding_latent=None, is_input_images:bool=False): - if isinstance(latents, list): - return_list = False - if padding_latent is None: - padding_latent = [None] * len(latents) - return_list = True - patched_latents, num_tokens, shapes = [], [], [] - for latent, padding in zip(latents, padding_latent): - height, width = latent.shape[-2:] - if is_input_images: - latent = self.input_x_embedder(latent) - else: - latent = self.x_embedder(latent) - pos_embed = self.cropped_pos_embed(height, width) - latent = latent + pos_embed - if padding is not None: - latent = torch.cat([latent, padding], dim=-2) - patched_latents.append(latent) - - num_tokens.append(pos_embed.size(1)) - shapes.append([height, width]) - if not return_list: - latents = torch.cat(patched_latents, dim=0) - else: - latents = patched_latents - else: - height, width = latents.shape[-2:] - if is_input_images: - latents = self.input_x_embedder(latents) - else: - latents = self.x_embedder(latents) - pos_embed = self.cropped_pos_embed(height, width) - latents = latents + pos_embed - num_tokens = latents.size(1) - shapes = [height, width] - return latents, num_tokens, shapes - - def forward(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, padding_latent=None, past_key_values=None, return_past_key_values=True): - input_is_list = isinstance(x, list) - x, num_tokens, shapes = self.patch_multiple_resolutions(x, padding_latent) - time_token = self.time_token(timestep, dtype=x[0].dtype).unsqueeze(1) - if input_img_latents is not None: - input_latents, _, _ = self.patch_multiple_resolutions(input_img_latents, is_input_images=True) - if input_ids is not None: - condition_embeds = self.llm.embed_tokens(input_ids).clone() - input_img_inx = 0 - for b_inx in input_image_sizes.keys(): - for start_inx, end_inx in input_image_sizes[b_inx]: - condition_embeds[b_inx, start_inx: end_inx] = input_latents[input_img_inx] - input_img_inx += 1 - if input_img_latents is not None: - assert input_img_inx == len(input_latents) - - input_emb = torch.cat([condition_embeds, time_token, x], dim=1) - else: - input_emb = torch.cat([time_token, x], dim=1) - output = self.llm(inputs_embeds=input_emb, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values) - output, past_key_values = output.last_hidden_state, output.past_key_values - if input_is_list: - image_embedding = output[:, -max(num_tokens):] - time_emb = self.t_embedder(timestep, dtype=x.dtype) - x = self.final_layer(image_embedding, time_emb) - latents = [] - for i in range(x.size(0)): - latent = x[i:i+1, :num_tokens[i]] - latent = self.unpatchify(latent, shapes[i][0], shapes[i][1]) - latents.append(latent) - else: - image_embedding = output[:, -num_tokens:] - time_emb = self.t_embedder(timestep, dtype=x.dtype) - x = self.final_layer(image_embedding, time_emb) - latents = self.unpatchify(x, shapes[0], shapes[1]) - - if return_past_key_values: - return latents, past_key_values - return latents - - @torch.no_grad() - def forward_with_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache): - """ - Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. - """ - self.llm.config.use_cache = use_kv_cache - model_out, past_key_values = self.forward(x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, past_key_values=past_key_values, return_past_key_values=True) - if use_img_cfg: - cond, uncond, img_cond = torch.split(model_out, len(model_out) // 3, dim=0) - cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond) - model_out = [cond, cond, cond] - else: - cond, uncond = torch.split(model_out, len(model_out) // 2, dim=0) - cond = uncond + cfg_scale * (cond - uncond) - model_out = [cond, cond] - return torch.cat(model_out, dim=0), past_key_values - - - @torch.no_grad() - def forward_with_separate_cfg(self, x, timestep, input_ids, input_img_latents, input_image_sizes, attention_mask, position_ids, cfg_scale, use_img_cfg, img_cfg_scale, past_key_values, use_kv_cache, return_past_key_values=True): - """ - Forward pass of DiT, but also batches the unconditional forward pass for classifier-free guidance. - """ - self.llm.config.use_cache = use_kv_cache - if past_key_values is None: - past_key_values = [None] * len(attention_mask) - - x = torch.split(x, len(x) // len(attention_mask), dim=0) - timestep = timestep.to(x[0].dtype) - timestep = torch.split(timestep, len(timestep) // len(input_ids), dim=0) - - model_out, pask_key_values = [], [] - for i in range(len(input_ids)): - temp_out, temp_pask_key_values = self.forward(x[i], timestep[i], input_ids[i], input_img_latents[i], input_image_sizes[i], attention_mask[i], position_ids[i], past_key_values[i]) - model_out.append(temp_out) - pask_key_values.append(temp_pask_key_values) - - if len(model_out) == 3: - cond, uncond, img_cond = model_out - cond = uncond + img_cfg_scale * (img_cond - uncond) + cfg_scale * (cond - img_cond) - model_out = [cond, cond, cond] - elif len(model_out) == 2: - cond, uncond = model_out - cond = uncond + cfg_scale * (cond - uncond) - model_out = [cond, cond] - else: - return model_out[0] - return torch.cat(model_out, dim=0), pask_key_values diff --git a/modules/omnigen/pipeline.py b/modules/omnigen/pipeline.py deleted file mode 100644 index a07467543..000000000 --- a/modules/omnigen/pipeline.py +++ /dev/null @@ -1,219 +0,0 @@ -import os -from typing import List, Union -from PIL import Image -import torch -from huggingface_hub import snapshot_download -from peft import PeftModel -from diffusers.models import AutoencoderKL -from diffusers.utils import replace_example_docstring -from .model import OmniGen -from .processor import OmniGenProcessor -from .scheduler import OmniGenScheduler - - -EXAMPLE_DOC_STRING = """ - Examples: - ```py - >>> from OmniGen import OmniGenPipeline - >>> pipe = FluxControlNetPipeline.from_pretrained( - ... base_model - ... ) - >>> prompt = "A woman holds a bouquet of flowers and faces the camera" - >>> image = pipe( - ... prompt, - ... guidance_scale=3.0, - ... num_inference_steps=50, - ... ).images[0] - >>> image.save("t2i.png") - ``` -""" - - -class OmniGenPipeline(): - def __init__( - self, - vae: AutoencoderKL, - model: OmniGen, - processor: OmniGenProcessor, - ): - super().__init__() - self.vae = vae - self.model = model - self.processor = processor - self.device = None - self.dtype: None - self.separate_cfg_infer: bool = True - self.use_kv_cache: bool = False - # omnigen does not inherit from diffusionpipeline so we hack it - self._internal_dict = { # pylint: disable=protected-access - 'vae': self.vae, - 'model': self.model, - 'processor': self.processor, - } - - @classmethod - def from_pretrained(cls, model_name, vae_path: str=None, cache_dir: str=None): - if not os.path.exists(model_name): - cache_dir = cache_dir or os.getenv('HF_HUB_CACHE') - model_name = snapshot_download(repo_id=model_name, - cache_dir=cache_dir, - ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5']) - model = OmniGen.from_pretrained(model_name) - processor = OmniGenProcessor.from_pretrained(model_name) - if os.path.exists(os.path.join(model_name, "vae")): - vae = AutoencoderKL.from_pretrained(os.path.join(model_name, "vae")) - else: - vae = AutoencoderKL.from_pretrained(vae_path or "stabilityai/sdxl-vae") - return cls(vae, model, processor) - - def merge_lora(self, lora_path: str): - model = PeftModel.from_pretrained(self.model, lora_path) - model.merge_and_unload() - self.model = model - - def to(self, device: Union[str, torch.device]): - if isinstance(device, str): - device = torch.device(device) - self.model.to(device) - self.vae.to(device) - - def vae_encode(self, x, dtype): - x = x.to(dtype) - if self.vae.config.shift_factor is not None: - x = self.vae.encode(x).latent_dist.sample() - x = (x - self.vae.config.shift_factor) * self.vae.config.scaling_factor - else: - x = self.vae.encode(x).latent_dist.sample().mul_(self.vae.config.scaling_factor) - x = x.to(dtype) - return x - - def move_to_device(self, data): - if isinstance(data, list): - return [x.to(self.device) for x in data] - return data.to(self.device) - - - @torch.no_grad() - @replace_example_docstring(EXAMPLE_DOC_STRING) - def __call__( - self, - prompt: Union[str, List[str]], - input_images: Union[List[str], List[List[str]]] = None, - height: int = 1024, - width: int = 1024, - num_inference_steps: int = 50, - guidance_scale: float = 3, - use_img_guidance: bool = True, - img_guidance_scale: float = 1.6, - output_type: str = 'latent', - seed: int = None, - ): - r""" - Function invoked when calling the pipeline for generation. - - Args: - prompt (`str` or `List[str]`): - The prompt or prompts to guide the image generation. - input_images (`List[str]` or `List[List[str]]`, *optional*): - The list of input images. We will replace the "<|image_i|>" in prompt with the 1-th image in list. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. The number must be a multiple of 16. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. The number must be a multiple of 16. - num_inference_steps (`int`, *optional*, defaults to 50): - The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. - guidance_scale (`float`, *optional*, defaults to 4.0): - Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598). - `guidance_scale` is defined as `w` of equation 2. of [Imagen - Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale > - 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, - usually at the expense of lower image quality. - use_img_guidance (`bool`, *optional*, defaults to True): - Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). - img_guidance_scale (`float`, *optional*, defaults to 1.6): - Defined as equation 3 in [Instrucpix2pix](https://arxiv.org/pdf/2211.09800). - self.separate_cfg_infer (`bool`, *optional*, defaults to False): - Perform inference on images with different guidance separately; this can save memory when generating images of large size at the expense of slower inference. - self.use_kv_cache (`bool`, *optional*, defaults to True): enable kv cache to speed up the inference - generator (`torch.Generator` or `List[torch.Generator]`, *optional*): - One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) - to make generation deterministic. - Examples: - - Returns: - A list with the generated images. - """ - assert height%16 == 0 and width%16 == 0 - if self.separate_cfg_infer: - self.use_kv_cache = False - # raise "Currently, don't support both self.use_kv_cache and self.separate_cfg_infer" - if input_images is None: - use_img_guidance = False - if isinstance(prompt, str): - prompt = [prompt] - input_images = [input_images] if input_images is not None else None - - input_data = self.processor(prompt, input_images, height=height, width=width, use_img_cfg=use_img_guidance, separate_cfg_input=self.separate_cfg_infer) - - num_prompt = len(prompt) - num_cfg = 2 if use_img_guidance else 1 - latent_size_h, latent_size_w = height//8, width//8 - - if seed is not None: - generator = torch.Generator(device=self.device).manual_seed(int(seed)) - else: - generator = None - latents = torch.randn(num_prompt, 4, latent_size_h, latent_size_w, device=self.device, generator=generator) - latents = torch.cat([latents]*(1+num_cfg), 0).to(self.dtype) - - input_img_latents = [] - if self.separate_cfg_infer: - for temp_pixel_values in input_data['input_pixel_values']: - temp_input_latents = [] - for img in temp_pixel_values: - img = self.vae_encode(img.to(self.device), self.dtype) - temp_input_latents.append(img) - input_img_latents.append(temp_input_latents) - else: - for img in input_data['input_pixel_values']: - img = self.vae_encode(img.to(self.device), self.dtype) - input_img_latents.append(img) - - model_kwargs = dict(input_ids=self.move_to_device(input_data['input_ids']), - input_img_latents=input_img_latents, - input_image_sizes=input_data['input_image_sizes'], - attention_mask=self.move_to_device(input_data["attention_mask"]), - position_ids=self.move_to_device(input_data["position_ids"]), - cfg_scale=guidance_scale, - img_cfg_scale=img_guidance_scale, - use_img_cfg=use_img_guidance, - use_kv_cache=self.use_kv_cache) - - if self.separate_cfg_infer: - func = self.model.forward_with_separate_cfg - else: - func = self.model.forward_with_cfg - self.model.to(self.dtype) - - scheduler = OmniGenScheduler(num_steps=num_inference_steps) - samples = scheduler(latents, func, model_kwargs, use_kv_cache=self.use_kv_cache) - samples = samples.chunk((1+num_cfg), dim=0)[0] - - if output_type == 'latent': - output_images = { 'images': samples } - return output_images - - samples = samples.to(self.vae.dtype) - if self.vae.config.shift_factor is not None: - samples = samples / self.vae.config.scaling_factor + self.vae.config.shift_factor - else: - samples = samples / self.vae.config.scaling_factor - samples = self.vae.decode(samples).sample - - output_samples = (samples * 0.5 + 0.5).clamp(0, 1)*255 - output_samples = output_samples.permute(0, 2, 3, 1).to("cpu", dtype=torch.uint8).numpy() - output_images = [] - for _i, sample in enumerate(output_samples): - output_images.append(Image.fromarray(sample)) - - return output_images diff --git a/modules/omnigen/processor.py b/modules/omnigen/processor.py deleted file mode 100644 index ada813a8b..000000000 --- a/modules/omnigen/processor.py +++ /dev/null @@ -1,312 +0,0 @@ -import os -import re -from typing import Dict, List -import torch -from torchvision import transforms -from transformers import AutoTokenizer -from huggingface_hub import snapshot_download -from .utils import crop_arr - - -class OmniGenProcessor: - def __init__(self, - text_tokenizer, - max_image_size: int=1024): - self.text_tokenizer = text_tokenizer - self.max_image_size = max_image_size - - self.image_transform = transforms.Compose([ - transforms.Lambda(lambda pil_image: crop_arr(pil_image, max_image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True) - ]) - - self.collator = OmniGenCollator() - self.separate_collator = OmniGenSeparateCollator() - - @classmethod - def from_pretrained(cls, model_name): - if not os.path.exists(model_name): - cache_folder = os.getenv('HF_HUB_CACHE') - model_name = snapshot_download(repo_id=model_name, - cache_dir=cache_folder, - allow_patterns="*.json") - text_tokenizer = AutoTokenizer.from_pretrained(model_name) - - return cls(text_tokenizer) - - - def process_image(self, image): - return self.image_transform(image) - - def process_multi_modal_prompt(self, text, input_images): - text = self.add_prefix_instruction(text) - if input_images is None or len(input_images) == 0: - model_inputs = self.text_tokenizer(text) - return {"input_ids": model_inputs.input_ids, "pixel_values": None, "image_sizes": None} - - pattern = r"<\|image_\d+\|>" - prompt_chunks = [self.text_tokenizer(chunk).input_ids for chunk in re.split(pattern, text)] - - for i in range(1, len(prompt_chunks)): - if prompt_chunks[i][0] == 1: - prompt_chunks[i] = prompt_chunks[i][1:] - - image_tags = re.findall(pattern, text) - image_ids = [int(s.split("|")[1].split("_")[-1]) for s in image_tags] - - unique_image_ids = sorted(list(set(image_ids))) - assert unique_image_ids == list(range(1, len(unique_image_ids)+1)), f"image_ids must start from 1, and must be continuous int, e.g. [1, 2, 3], cannot be {unique_image_ids}" - # total images must be the same as the number of image tags - assert len(unique_image_ids) == len(input_images), f"total images must be the same as the number of image tags, got {len(unique_image_ids)} image tags and {len(input_images)} images" - - input_images = [input_images[x-1] for x in image_ids] - - all_input_ids = [] - img_inx = [] - _idx = 0 - for i in range(len(prompt_chunks)): - all_input_ids.extend(prompt_chunks[i]) - if i != len(prompt_chunks) -1: - start_inx = len(all_input_ids) - size = input_images[i].size(-2) * input_images[i].size(-1) // 16 // 16 - img_inx.append([start_inx, start_inx+size]) - all_input_ids.extend([0]*size) - - return {"input_ids": all_input_ids, "pixel_values": input_images, "image_sizes": img_inx} - - - def add_prefix_instruction(self, prompt): - user_prompt = '<|user|>\n' - generation_prompt = 'Generate an image according to the following instructions\n' - assistant_prompt = '<|assistant|>\n<|diffusion|>' - prompt_suffix = "<|end|>\n" - prompt = f"{user_prompt}{generation_prompt}{prompt}{prompt_suffix}{assistant_prompt}" - return prompt - - - def __call__(self, - instructions: List[str], - input_images: List[List[str]] = None, - height: int = 1024, - width: int = 1024, - negative_prompt: str = "low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers.", - use_img_cfg: bool = True, - separate_cfg_input: bool = False, - ) -> Dict: - - if input_images is None: - use_img_cfg = False - if isinstance(instructions, str): - instructions = [instructions] - input_images = [input_images] - - input_data = [] - for i in range(len(instructions)): - cur_instruction = instructions[i] - cur_input_images = None if input_images is None else input_images[i] - if cur_input_images is not None and len(cur_input_images) > 0: - cur_input_images = [self.process_image(x) for x in cur_input_images] - else: - cur_input_images = None - assert "<|image_1|>" not in cur_instruction - - mllm_input = self.process_multi_modal_prompt(cur_instruction, cur_input_images) - - neg_mllm_input, img_cfg_mllm_input = None, None - neg_mllm_input = self.process_multi_modal_prompt(negative_prompt, None) - if use_img_cfg: - if cur_input_images is not None and len(cur_input_images) >= 1: - img_cfg_prompt = [f"<|image_{i+1}|>" for i in range(len(cur_input_images))] - img_cfg_mllm_input = self.process_multi_modal_prompt(" ".join(img_cfg_prompt), cur_input_images) - else: - img_cfg_mllm_input = neg_mllm_input - - input_data.append((mllm_input, neg_mllm_input, img_cfg_mllm_input, [height, width])) - - if separate_cfg_input: - return self.separate_collator(input_data) - return self.collator(input_data) - - -class OmniGenCollator: - def __init__(self, pad_token_id=2, hidden_size=3072): - self.pad_token_id = pad_token_id - self.hidden_size = hidden_size - - def create_position(self, attention_mask, num_tokens_for_output_images): - position_ids = [] - text_length = attention_mask.size(-1) - img_length = max(num_tokens_for_output_images) - for mask in attention_mask: - temp_l = torch.sum(mask) - temp_position = [0]*(text_length-temp_l) + [i for i in range(temp_l+img_length+1)] # we add a time embedding into the sequence, so add one more token - position_ids.append(temp_position) - return torch.LongTensor(position_ids) - - def create_mask(self, attention_mask, num_tokens_for_output_images): - extended_mask = [] - padding_images = [] - text_length = attention_mask.size(-1) - img_length = max(num_tokens_for_output_images) - seq_len = text_length + img_length + 1 # we add a time embedding into the sequence, so add one more token - inx = 0 - for mask in attention_mask: - temp_l = torch.sum(mask) - pad_l = text_length - temp_l - - temp_mask = torch.tril(torch.ones(size=(temp_l+1, temp_l+1))) - - image_mask = torch.zeros(size=(temp_l+1, img_length)) - temp_mask = torch.cat([temp_mask, image_mask], dim=-1) - - image_mask = torch.ones(size=(img_length, temp_l+img_length+1)) - temp_mask = torch.cat([temp_mask, image_mask], dim=0) - - if pad_l > 0: - pad_mask = torch.zeros(size=(temp_l+1+img_length, pad_l)) - temp_mask = torch.cat([pad_mask, temp_mask], dim=-1) - - pad_mask = torch.ones(size=(pad_l, seq_len)) - temp_mask = torch.cat([pad_mask, temp_mask], dim=0) - - true_img_length = num_tokens_for_output_images[inx] - pad_img_length = img_length - true_img_length - if pad_img_length > 0: - temp_mask[:, -pad_img_length:] = 0 - temp_padding_imgs = torch.zeros(size=(1, pad_img_length, self.hidden_size)) - else: - temp_padding_imgs = None - - extended_mask.append(temp_mask.unsqueeze(0)) - padding_images.append(temp_padding_imgs) - inx += 1 - return torch.cat(extended_mask, dim=0), padding_images - - def adjust_attention_for_input_images(self, attention_mask, image_sizes): - for b_inx in image_sizes.keys(): - for start_inx, end_inx in image_sizes[b_inx]: - attention_mask[b_inx][start_inx:end_inx, start_inx:end_inx] = 1 - - return attention_mask - - def pad_input_ids(self, input_ids, image_sizes): - max_l = max([len(x) for x in input_ids]) - padded_ids = [] - attention_mask = [] - _new_image_sizes = [] - - for i in range(len(input_ids)): - temp_ids = input_ids[i] - temp_l = len(temp_ids) - pad_l = max_l - temp_l - if pad_l == 0: - attention_mask.append([1]*max_l) - padded_ids.append(temp_ids) - else: - attention_mask.append([0]*pad_l+[1]*temp_l) - padded_ids.append([self.pad_token_id]*pad_l+temp_ids) - - if i in image_sizes: - new_inx = [] - for old_inx in image_sizes[i]: - new_inx.append([x+pad_l for x in old_inx]) - image_sizes[i] = new_inx - - return torch.LongTensor(padded_ids), torch.LongTensor(attention_mask), image_sizes - - - def process_mllm_input(self, mllm_inputs, target_img_size): - num_tokens_for_output_images = [] - for img_size in target_img_size: - num_tokens_for_output_images.append(img_size[0]*img_size[1]//16//16) - - pixel_values, image_sizes = [], {} - b_inx = 0 - for x in mllm_inputs: - if x['pixel_values'] is not None: - pixel_values.extend(x['pixel_values']) - for size in x['image_sizes']: - if b_inx not in image_sizes: - image_sizes[b_inx] = [size] - else: - image_sizes[b_inx].append(size) - b_inx += 1 - pixel_values = [x.unsqueeze(0) for x in pixel_values] - - input_ids = [x['input_ids'] for x in mllm_inputs] - padded_input_ids, attention_mask, image_sizes = self.pad_input_ids(input_ids, image_sizes) - position_ids = self.create_position(attention_mask, num_tokens_for_output_images) - attention_mask, padding_images = self.create_mask(attention_mask, num_tokens_for_output_images) - attention_mask = self.adjust_attention_for_input_images(attention_mask, image_sizes) - - return padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes - - def __call__(self, features): - mllm_inputs = [f[0] for f in features] - cfg_mllm_inputs = [f[1] for f in features] - img_cfg_mllm_input = [f[2] for f in features] - target_img_size = [f[3] for f in features] - - if img_cfg_mllm_input[0] is not None: - mllm_inputs = mllm_inputs + cfg_mllm_inputs + img_cfg_mllm_input - target_img_size = target_img_size + target_img_size + target_img_size - else: - mllm_inputs = mllm_inputs + cfg_mllm_inputs - target_img_size = target_img_size + target_img_size - - - all_padded_input_ids, all_position_ids, all_attention_mask, all_padding_images, all_pixel_values, all_image_sizes = self.process_mllm_input(mllm_inputs, target_img_size) - - data = {"input_ids": all_padded_input_ids, - "attention_mask": all_attention_mask, - "position_ids": all_position_ids, - "input_pixel_values": all_pixel_values, - "input_image_sizes": all_image_sizes, - "padding_images": all_padding_images, - } - return data - - -class OmniGenSeparateCollator(OmniGenCollator): - def __call__(self, features): - mllm_inputs = [f[0] for f in features] - cfg_mllm_inputs = [f[1] for f in features] - img_cfg_mllm_input = [f[2] for f in features] - target_img_size = [f[3] for f in features] - - all_padded_input_ids, all_attention_mask, all_position_ids, all_pixel_values, all_image_sizes, all_padding_images = [], [], [], [], [], [] - - padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(mllm_inputs, target_img_size) - all_padded_input_ids.append(padded_input_ids) - all_attention_mask.append(attention_mask) - all_position_ids.append(position_ids) - all_pixel_values.append(pixel_values) - all_image_sizes.append(image_sizes) - all_padding_images.append(padding_images) - - if cfg_mllm_inputs[0] is not None: - padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(cfg_mllm_inputs, target_img_size) - all_padded_input_ids.append(padded_input_ids) - all_attention_mask.append(attention_mask) - all_position_ids.append(position_ids) - all_pixel_values.append(pixel_values) - all_image_sizes.append(image_sizes) - all_padding_images.append(padding_images) - if img_cfg_mllm_input[0] is not None: - padded_input_ids, position_ids, attention_mask, padding_images, pixel_values, image_sizes = self.process_mllm_input(img_cfg_mllm_input, target_img_size) - all_padded_input_ids.append(padded_input_ids) - all_attention_mask.append(attention_mask) - all_position_ids.append(position_ids) - all_pixel_values.append(pixel_values) - all_image_sizes.append(image_sizes) - all_padding_images.append(padding_images) - - data = {"input_ids": all_padded_input_ids, - "attention_mask": all_attention_mask, - "position_ids": all_position_ids, - "input_pixel_values": all_pixel_values, - "input_image_sizes": all_image_sizes, - "padding_images": all_padding_images, - } - return data diff --git a/modules/omnigen/scheduler.py b/modules/omnigen/scheduler.py deleted file mode 100644 index 0764fd8f0..000000000 --- a/modules/omnigen/scheduler.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch -from tqdm import tqdm -from transformers.cache_utils import Cache, DynamicCache - -class OmniGenScheduler: - def __init__(self, num_steps: int=50, time_shifting_factor: int=1): - self.num_steps = num_steps - self.time_shift = time_shifting_factor - - t = torch.linspace(0, 1, num_steps+1) - t = t / (t + time_shifting_factor - time_shifting_factor * t) - self.sigma = t - - def crop_kv_cache(self, past_key_values, num_tokens_for_img): - crop_past_key_values = () - for layer_idx in range(len(past_key_values)): - key_states, value_states = past_key_values[layer_idx][:2] - crop_past_key_values += ((key_states[..., :-(num_tokens_for_img+1), :], value_states[..., :-(num_tokens_for_img+1), :], ),) - return crop_past_key_values - # return DynamicCache.from_legacy_cache(crop_past_key_values) - - def crop_position_ids_for_cache(self, position_ids, num_tokens_for_img): - if isinstance(position_ids, list): - for i in range(len(position_ids)): - position_ids[i] = position_ids[i][:, -(num_tokens_for_img+1):] - else: - position_ids = position_ids[:, -(num_tokens_for_img+1):] - return position_ids - - def crop_attention_mask_for_cache(self, attention_mask, num_tokens_for_img): - if isinstance(attention_mask, list): - return [x[..., -(num_tokens_for_img+1):, :] for x in attention_mask] - return attention_mask[..., -(num_tokens_for_img+1):, :] - - def __call__(self, z, func, model_kwargs, use_kv_cache: bool=True): - past_key_values = None - for i in tqdm(range(self.num_steps)): - timesteps = torch.zeros(size=(len(z), )).to(z.device) + self.sigma[i] - pred, temp_past_key_values = func(z, timesteps, past_key_values=past_key_values, **model_kwargs) - sigma_next = self.sigma[i+1] - sigma = self.sigma[i] - z = z + (sigma_next - sigma) * pred - if i == 0 and use_kv_cache: - num_tokens_for_img = z.size(-1)*z.size(-2) // 4 - if isinstance(temp_past_key_values, list): - past_key_values = [self.crop_kv_cache(x, num_tokens_for_img) for x in temp_past_key_values] - model_kwargs['input_ids'] = [None] * len(temp_past_key_values) - else: - past_key_values = self.crop_kv_cache(temp_past_key_values, num_tokens_for_img) - model_kwargs['input_ids'] = None - - model_kwargs['position_ids'] = self.crop_position_ids_for_cache(model_kwargs['position_ids'], num_tokens_for_img) - model_kwargs['attention_mask'] = self.crop_attention_mask_for_cache(model_kwargs['attention_mask'], num_tokens_for_img) - return z - diff --git a/modules/omnigen/transformer.py b/modules/omnigen/transformer.py deleted file mode 100644 index d166309ca..000000000 --- a/modules/omnigen/transformer.py +++ /dev/null @@ -1,164 +0,0 @@ -import math -import warnings -from typing import List, Optional, Tuple, Union - -import torch -import torch.utils.checkpoint -from torch import nn -from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss -from huggingface_hub import snapshot_download - -from transformers.modeling_outputs import ( - BaseModelOutputWithPast, - CausalLMOutputWithPast, - SequenceClassifierOutputWithPast, - TokenClassifierOutput, -) -from transformers.modeling_utils import PreTrainedModel -from transformers import Phi3Config, Phi3Model -from transformers.cache_utils import Cache, DynamicCache, StaticCache -from transformers.utils import logging - -logger = logging.get_logger(__name__) - - -class Phi3Transformer(Phi3Model): - """ - Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`] - We only modified the attention mask - Args: - config: Phi3Config - """ - - def forward( - self, - input_ids: torch.LongTensor = None, - attention_mask: Optional[torch.Tensor] = None, - position_ids: Optional[torch.LongTensor] = None, - past_key_values: Optional[List[torch.FloatTensor]] = None, - inputs_embeds: Optional[torch.FloatTensor] = None, - use_cache: Optional[bool] = None, - output_attentions: Optional[bool] = None, - output_hidden_states: Optional[bool] = None, - return_dict: Optional[bool] = None, - cache_position: Optional[torch.LongTensor] = None, - ) -> Union[Tuple, BaseModelOutputWithPast]: - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions - output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states - ) - use_cache = use_cache if use_cache is not None else self.config.use_cache - - return_dict = return_dict if return_dict is not None else self.config.use_return_dict - - if (input_ids is None) ^ (inputs_embeds is not None): - raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - - if self.gradient_checkpointing and self.training: - if use_cache: - logger.warning_once( - "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." - ) - use_cache = False - - # kept for BC (non `Cache` `past_key_values` inputs) - return_legacy_cache = False - if use_cache and not isinstance(past_key_values, Cache): - return_legacy_cache = True - if past_key_values is None: - past_key_values = DynamicCache() - else: - past_key_values = DynamicCache.from_legacy_cache(past_key_values) - logger.warning_once( - "We detected that you are passing `past_key_values` as a tuple of tuples. This is deprecated and " - "will be removed in v4.47. Please convert your cache or use an appropriate `Cache` class " - "(https://huggingface.co/docs/transformers/kv_cache#legacy-cache-format)" - ) - - if inputs_embeds is None: - inputs_embeds = self.embed_tokens(input_ids) - - if cache_position is None: - past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 - cache_position = torch.arange( - past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device - ) - if position_ids is None: - position_ids = cache_position.unsqueeze(0) - - if attention_mask is not None and attention_mask.dim() == 3: - dtype = inputs_embeds.dtype - min_dtype = torch.finfo(dtype).min - attention_mask = (1 - attention_mask) * min_dtype - attention_mask = attention_mask.unsqueeze(1).to(inputs_embeds.dtype) - else: - raise - # causal_mask = self._update_causal_mask( - # attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions - # ) - - hidden_states = inputs_embeds - - # create position embeddings to be shared across the decoder layers - position_embeddings = self.rotary_emb(hidden_states, position_ids) - - # decoder layers - all_hidden_states = () if output_hidden_states else None - all_self_attns = () if output_attentions else None - next_decoder_cache = None - - for decoder_layer in self.layers: - if output_hidden_states: - all_hidden_states += (hidden_states,) - - if self.gradient_checkpointing and self.training: - layer_outputs = self._gradient_checkpointing_func( - decoder_layer.__call__, - hidden_states, - attention_mask, - position_ids, - past_key_values, - output_attentions, - use_cache, - cache_position, - position_embeddings, - ) - else: - layer_outputs = decoder_layer( - hidden_states, - attention_mask=attention_mask, - position_ids=position_ids, - past_key_value=past_key_values, - output_attentions=output_attentions, - use_cache=use_cache, - cache_position=cache_position, - position_embeddings=position_embeddings, - ) - - hidden_states = layer_outputs[0] - - if use_cache: - next_decoder_cache = layer_outputs[2 if output_attentions else 1] - - if output_attentions: - all_self_attns += (layer_outputs[1],) - - hidden_states = self.norm(hidden_states) - - # add hidden states from the last decoder layer - if output_hidden_states: - all_hidden_states += (hidden_states,) - - next_cache = next_decoder_cache if use_cache else None - if return_legacy_cache: - next_cache = next_cache.to_legacy_cache() - - if not return_dict: - return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None) - return BaseModelOutputWithPast( - last_hidden_state=hidden_states, - past_key_values=next_cache, - hidden_states=all_hidden_states, - attentions=all_self_attns, - ) - diff --git a/modules/omnigen/utils.py b/modules/omnigen/utils.py deleted file mode 100644 index bf0a6de62..000000000 --- a/modules/omnigen/utils.py +++ /dev/null @@ -1,105 +0,0 @@ -import logging - -from PIL import Image -import torch -import numpy as np - -def create_logger(logging_dir): - """ - Create a logger that writes to a log file and stdout. - """ - logging.basicConfig( - level=logging.INFO, - format='[\033[34m%(asctime)s\033[0m] %(message)s', - datefmt='%Y-%m-%d %H:%M:%S', - handlers=[logging.StreamHandler(), logging.FileHandler(f"{logging_dir}/log.txt")] - ) - logger = logging.getLogger(__name__) - return logger - - -@torch.no_grad() -def update_ema(ema_model, model, decay=0.9999): - """ - Step the EMA model towards the current model. - """ - ema_params = dict(ema_model.named_parameters()) - for name, param in model.named_parameters(): - ema_params[name].mul_(decay).add_(param.data, alpha=1 - decay) - - -def requires_grad(model, flag=True): - """ - Set requires_grad flag for all parameters in a model. - """ - for p in model.parameters(): - p.requires_grad = flag - - -def center_crop_arr(pil_image, image_size): - """ - Center cropping implementation from ADM. - https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 - """ - while min(*pil_image.size) >= 2 * image_size: - pil_image = pil_image.resize( - tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS - ) - - scale = image_size / min(*pil_image.size) - pil_image = pil_image.resize( - tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS - ) - - arr = np.array(pil_image) - crop_y = (arr.shape[0] - image_size) // 2 - crop_x = (arr.shape[1] - image_size) // 2 - return Image.fromarray(arr[crop_y: crop_y + image_size, crop_x: crop_x + image_size]) - - -def crop_arr(pil_image, max_image_size): - while min(*pil_image.size) >= 2 * max_image_size: - pil_image = pil_image.resize( - tuple(x // 2 for x in pil_image.size), resample=Image.Resampling.LANCZOS - ) - - if max(*pil_image.size) > max_image_size: - scale = max_image_size / max(*pil_image.size) - pil_image = pil_image.resize( - tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS - ) - - if min(*pil_image.size) < 16: - scale = 16 / min(*pil_image.size) - pil_image = pil_image.resize( - tuple(round(x * scale) for x in pil_image.size), resample=Image.Resampling.LANCZOS - ) - - arr = np.array(pil_image) - crop_y1 = (arr.shape[0] % 16) // 2 - crop_y2 = arr.shape[0] % 16 - crop_y1 - - crop_x1 = (arr.shape[1] % 16) // 2 - crop_x2 = arr.shape[1] % 16 - crop_x1 - - arr = arr[crop_y1:arr.shape[0]-crop_y2, crop_x1:arr.shape[1]-crop_x2] - return Image.fromarray(arr) - - -def vae_encode(vae, x, weight_dtype): - if x is not None: - if vae.config.shift_factor is not None: - x = vae.encode(x).latent_dist.sample() - x = (x - vae.config.shift_factor) * vae.config.scaling_factor - else: - x = vae.encode(x).latent_dist.sample().mul_(vae.config.scaling_factor) - x = x.to(weight_dtype) - return x - - -def vae_encode_list(vae, x, weight_dtype): - latents = [] - for img in x: - img = vae_encode(vae, img, weight_dtype) - latents.append(img) - return latents diff --git a/modules/processing_args.py b/modules/processing_args.py index 6ef43667b..d9ad9869f 100644 --- a/modules/processing_args.py +++ b/modules/processing_args.py @@ -166,7 +166,7 @@ def set_pipeline_args(p, model, prompts:list, negative_prompts:list, prompts_2:t extra_networks.activate(p, include=['text_encoder', 'text_encoder_2', 'text_encoder_3']) if 'prompt' in possible: if 'OmniGen' in model.__class__.__name__: - prompts = [p.replace('|image|', '<|image_1|>') for p in prompts] + prompts = [p.replace('|image|', '<|image_1|>') for p in prompts] if 'HiDreamImage' in model.__class__.__name__ and prompt_parser_diffusers.embedder is not None: args['pooled_prompt_embeds'] = prompt_parser_diffusers.embedder('positive_pooleds') prompt_embeds = prompt_parser_diffusers.embedder('prompt_embeds') diff --git a/modules/sd_offload.py b/modules/sd_offload.py index c57165fb5..57057d233 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -14,7 +14,7 @@ debug_move = log.trace if debug else lambda *args, **kwargs: None offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'cogview4'] offload_post = ['h1'] offload_hook_instance = None -balanced_offload_exclude = ['OmniGenPipeline', 'CogView4Pipeline'] +balanced_offload_exclude = ['CogView4Pipeline'] def get_signature(cls): diff --git a/modules/sd_vae_remote.py b/modules/sd_vae_remote.py index 741d349bc..099c48148 100644 --- a/modules/sd_vae_remote.py +++ b/modules/sd_vae_remote.py @@ -12,14 +12,25 @@ hf_decode_endpoints = { 'sd': 'https://q1bj3bpq6kzilnsu.us-east-1.aws.endpoints.huggingface.cloud', 'sdxl': 'https://x2dmsqunjd6k9prw.us-east-1.aws.endpoints.huggingface.cloud', 'f1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud', - 'h1': 'https://whhx50ex1aryqvw6.us-east-1.aws.endpoints.huggingface.cloud', 'hunyuanvideo': 'https://o7ywnmrahorts457.us-east-1.aws.endpoints.huggingface.cloud', } +hf_decode_endpoints['pixartalpha'] = hf_decode_endpoints['sd'] +hf_decode_endpoints['pixartsigma'] = hf_decode_endpoints['sdxl'] +hf_decode_endpoints['omnigen'] = hf_decode_endpoints['sdxl'] +hf_decode_endpoints['h1'] = hf_decode_endpoints['f1'] +hf_decode_endpoints['lumina2'] = hf_decode_endpoints['f1'] + hf_encode_endpoints = { 'sd': 'https://qc6479g0aac6qwy9.us-east-1.aws.endpoints.huggingface.cloud', 'sdxl': 'https://xjqqhmyn62rog84g.us-east-1.aws.endpoints.huggingface.cloud', 'f1': 'https://ptccx55jz97f9zgo.us-east-1.aws.endpoints.huggingface.cloud', } +hf_encode_endpoints['pixartalpha'] = hf_encode_endpoints['sd'] +hf_encode_endpoints['pixartsigma'] = hf_encode_endpoints['sdxl'] +hf_encode_endpoints['omnigen'] = hf_encode_endpoints['sdxl'] +hf_encode_endpoints['h1'] = hf_encode_endpoints['f1'] +hf_encode_endpoints['lumina2'] = hf_encode_endpoints['f1'] + dtypes = { "float16": torch.float16, "float32": torch.float32, @@ -74,7 +85,7 @@ def remote_decode(latents: torch.Tensor, width: int = 0, height: int = 0, model_ params["output_type"] = "pt" params["output_tensor_type"] = "binary" headers["Accept"] = "tensor/binary" - if (model_type == 'f1' or model_type == 'h1') and (width > 0) and (height > 0): + if model_type in {'f1', 'h1', 'lumina2'} and (width > 0) and (height > 0): params['width'] = width params['height'] = height if shared.sd_model.vae is not None and shared.sd_model.vae.config is not None: diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 3bcd1b322..7837c2b14 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -36,7 +36,7 @@ prev_cls = '' prev_type = '' prev_model = '' lock = threading.Lock() -supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha'] +supported = ['sd', 'sdxl', 'f1', 'h1', 'lumina2', 'hunyuanvideo', 'wanvideo', 'mochivideo', 'pixartsigma', 'pixartalpha', 'omnigen'] def warn_once(msg, variant=None): @@ -57,7 +57,7 @@ def get_model(model_type = 'decoder', variant = None): cls = 'sd' elif cls in {'h1', 'lumina2'}: cls = 'f1' - elif cls == 'pixartsigma': + elif cls in {'pixartsigma', 'omnigen'}: cls = 'sdxl' elif cls not in supported: warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant) diff --git a/modules/shared_items.py b/modules/shared_items.py index 9a09b8c7f..cdc1f1980 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -41,10 +41,10 @@ pipelines = { 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), 'Amused': getattr(diffusers, 'AmusedPipeline', None), 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), + 'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), # dynamically imported and redefined later 'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser - 'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser 'InstaFlow': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser 'SegMoE': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser } From 77db759d88ee23e498d8afe51326986020d302a0 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Tue, 24 Jun 2025 13:15:31 +0300 Subject: [PATCH 11/25] Cleanup --- modules/shared_items.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/modules/shared_items.py b/modules/shared_items.py index cdc1f1980..baf247a35 100644 --- a/modules/shared_items.py +++ b/modules/shared_items.py @@ -41,7 +41,7 @@ pipelines = { 'UniDiffuser': getattr(diffusers, 'UniDiffuserPipeline', None), 'Amused': getattr(diffusers, 'AmusedPipeline', None), 'HiDream': getattr(diffusers, 'HiDreamImagePipeline', None), - 'OmniGenPipeline': getattr(diffusers, 'DiffusionPipeline', None), + 'OmniGenPipeline': getattr(diffusers, 'OmniGenPipeline', None), # dynamically imported and redefined later 'Meissonic': getattr(diffusers, 'DiffusionPipeline', None), # dynamically redefined and loaded in sd_models.load_diffuser From b650618b274b674588155e7c8d856c6d85969a71 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 09:11:17 -0400 Subject: [PATCH 12/25] fix params.txt Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 8 +++++--- javascript/gallery.js | 2 +- modules/generation_parameters_copypaste.py | 7 +++---- modules/gr_tempdir.py | 5 +++-- modules/images.py | 6 +++--- modules/paths.py | 1 + modules/ui_extra_networks.py | 10 ++++------ 7 files changed, 20 insertions(+), 19 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5faadfa04..25985d83b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,10 +1,10 @@ # Change Log for SD.Next -## Update for 2025-06-24 +## Update for 2025-06-25 - **Changes** - - Use Diffusers version of OmniGen - - Support Remote VAE with Omnigen, Lumina 2 and PixArt + - Use Diffusers version of *OmniGen* + - Support Remote VAE with *Omnigen, Lumina 2 and PixArt* - **SDNQ Quantization** - Add modules_to_not_convert support for post mode @@ -18,6 +18,8 @@ - LTXVideo default scheduler - Balanced offload with OmniGen - Quantization with OmniGen + - Do not save empty `params.txt` file + - Override `params.txt` using `SD_PATH_PARAMS` env variable ## Update for 2025-06-16 diff --git a/javascript/gallery.js b/javascript/gallery.js index 32ae0bde4..54945cf97 100644 --- a/javascript/gallery.js +++ b/javascript/gallery.js @@ -286,7 +286,7 @@ async function gallerySearch(evt) { const findDuplicates = (arr, key) => { const map = new Map(); - return arr.filter(item => { + return arr.filter((item) => { const value = item[key]; if (map.has(value)) return true; map.set(value, true); diff --git a/modules/generation_parameters_copypaste.py b/modules/generation_parameters_copypaste.py index 98c2b39cd..6393c9051 100644 --- a/modules/generation_parameters_copypaste.py +++ b/modules/generation_parameters_copypaste.py @@ -3,7 +3,7 @@ import io import os from PIL import Image import gradio as gr -from modules.paths import data_path +from modules.paths import params_path from modules import shared, gr_tempdir, script_callbacks, images from modules.infotext import parse, mapping, quote, unquote # pylint: disable=unused-import @@ -223,9 +223,8 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp def paste_func(prompt): if prompt is None or len(prompt.strip()) == 0: - filename = os.path.join(data_path, "params.txt") - if os.path.exists(filename): - with open(filename, "r", encoding="utf8") as file: + if os.path.exists(params_path): + with open(params_path, "r", encoding="utf8") as file: prompt = file.read() shared.log.debug(f'Prompt parse: type="params" prompt="{prompt}"') else: diff --git a/modules/gr_tempdir.py b/modules/gr_tempdir.py index bbe2b2192..f19fd53a3 100644 --- a/modules/gr_tempdir.py +++ b/modules/gr_tempdir.py @@ -74,8 +74,9 @@ def pil_to_temp_file(self, img: Image, dir: str, format="png") -> str: # pylint: shared.state.image_history += 1 params = ', '.join([f'{k}: {v}' for k, v in img.info.items()]) params = params[12:] if params.startswith('parameters: ') else params - with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file: - file.write(params) + if len(params) > 2: + with open(paths.params_path, "w", encoding="utf8") as file: + file.write(params) return name diff --git a/modules/images.py b/modules/images.py index bd3cd0c70..e8a0fa384 100644 --- a/modules/images.py +++ b/modules/images.py @@ -26,7 +26,6 @@ except Exception: pass - def sanitize_filename_part(text, replace_spaces=True): if text is None: return None @@ -47,8 +46,9 @@ def atomically_save_image(): while True: image, filename, extension, params, exifinfo, filename_txt = save_queue.get() shared.state.image_history += 1 - with open(os.path.join(paths.data_path, "params.txt"), "w", encoding="utf8") as file: - file.write(exifinfo) + if len(exifinfo) > 2: + with open(paths.params_path, "w", encoding="utf8") as file: + file.write(exifinfo) fn = filename + extension filename = filename.strip() if extension[0] != '.': # add dot if missing diff --git a/modules/paths.py b/modules/paths.py index 06cf34bff..d3257f043 100644 --- a/modules/paths.py +++ b/modules/paths.py @@ -29,6 +29,7 @@ script_path = os.path.dirname(modules_path) data_path = cli.data_dir models_config = cli.models_dir or config.get('models_dir') or 'models' models_path = models_config if os.path.isabs(models_config) else os.path.join(data_path, models_config) +params_path = os.environ.get('SD_PATH_PARAMS', os.path.join(data_path, "params.txt")) extensions_dir = cli.extensions_dir or os.path.join(data_path, "extensions") extensions_builtin_dir = "extensions-builtin" sd_configs_path = os.path.join(script_path, "configs") diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py index d2f8e4576..ca4c5c855 100644 --- a/modules/ui_extra_networks.py +++ b/modules/ui_extra_networks.py @@ -941,9 +941,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): from modules.processing_info import get_last_args params, text = get_last_args() if (not params) or (not text) or (len(text) == 0): - filename = os.path.join(paths.data_path, "params.txt") - if os.path.exists(filename): - with open(filename, "r", encoding="utf8") as file: + if os.path.exists(paths.params_path): + with open(paths.params_path, "r", encoding="utf8") as file: text = file.read() else: text = '' @@ -960,9 +959,8 @@ def create_ui(container, button_parent, tabname, skip_indexing = False): from modules.processing_info import get_last_args params, text = get_last_args() if (not params) or (not text) or (len(text) == 0): - fn = os.path.join(paths.data_path, "params.txt") - if os.path.exists(fn): - with open(fn, "r", encoding="utf8") as file: + if os.path.exists(paths.params_path): + with open(paths.params_path, "r", encoding="utf8") as file: text = file.read() else: text = '' From 82e5848259701bcc08e2fe1a35da46313871f3a7 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 09:24:33 -0400 Subject: [PATCH 13/25] add wheel to requirements Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + requirements.txt | 1 + 2 files changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 25985d83b..a01e5e35e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -20,6 +20,7 @@ - Quantization with OmniGen - Do not save empty `params.txt` file - Override `params.txt` using `SD_PATH_PARAMS` env variable + - Add `wheel` to requirements due to `pip` change ## Update for 2025-06-16 diff --git a/requirements.txt b/requirements.txt index f5eb92014..5d88bbf35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ # required for python 3.12 setuptools==69.5.1 +wheel # standard patch-ng From a5d784f073a8bdc561a4f2f13b6a1b758238ab7d Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 10:30:59 -0400 Subject: [PATCH 14/25] fix delete file from main gallery view Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ javascript/ui.js | 11 +++++++++ modules/processing_helpers.py | 5 +--- modules/sd_samplers.py | 11 +++++++++ modules/ui_common.py | 43 ++++++++++++++++++----------------- 5 files changed, 47 insertions(+), 25 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a01e5e35e..34d357111 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -21,6 +21,8 @@ - Do not save empty `params.txt` file - Override `params.txt` using `SD_PATH_PARAMS` env variable - Add `wheel` to requirements due to `pip` change + - Case-insensitive sampler name matching + - Fix delete file with gallery views ## Update for 2025-06-16 diff --git a/javascript/ui.js b/javascript/ui.js index b8112d523..ea8af5902 100644 --- a/javascript/ui.js +++ b/javascript/ui.js @@ -60,6 +60,17 @@ function selected_gallery_index() { return result; } +function selected_gallery_files() { + let allImages = []; + try { + let allCurrentButtons = gradioApp().querySelectorAll('[style="display: block;"].tabitem div[id$=_gallery].gradio-gallery .thumbnail-item.thumbnail-small'); + if (allCurrentButtons.length === 0) allCurrentButtons = gradioApp().querySelectorAll('.gradio-gallery .thumbnails > .thumbnail-item.thumbnail-small'); + allImages = Array.from(allCurrentButtons).map((v) => v.querySelector('img')?.src); + } catch { /**/ } + const selectedIndex = selected_gallery_index(); + return [allImages, selectedIndex]; +} + function extract_image_from_gallery(gallery) { if (gallery.length === 0) return [null]; if (gallery.length === 1) return [gallery[0]]; diff --git a/modules/processing_helpers.py b/modules/processing_helpers.py index 0f2d7bc6c..461e0e9ed 100644 --- a/modules/processing_helpers.py +++ b/modules/processing_helpers.py @@ -572,10 +572,7 @@ def update_sampler(p, sd_model, second_pass=False): if hasattr(sd_model, 'scheduler'): if sampler_selection == 'None': return - if sampler_selection is None: - sampler = sd_samplers.all_samplers_map.get("UniPC") - else: - sampler = sd_samplers.all_samplers_map.get(sampler_selection, None) + sampler = sd_samplers.find_sampler(sampler_selection) if sampler is None: shared.log.warning(f'Sampler: sampler="{sampler_selection}" not found') sampler = sd_samplers.all_samplers_map.get("UniPC") diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py index f73520027..16452da23 100644 --- a/modules/sd_samplers.py +++ b/modules/sd_samplers.py @@ -16,6 +16,17 @@ flow_models = ['Flux', 'StableDiffusion3', 'Lumina', 'AuraFlow', 'Sana', 'CogVi flow_models += ['Hunyuan', 'LTX', 'Mochi'] +def find_sampler(name:str): + if name is None or name == 'None': + return all_samplers_map.get("UniPC", None) + for sampler in all_samplers: + if sampler.name.lower() == name.lower() or name in sampler.aliases: + debug(f'Find sampler: name="{name}" found={sampler.name}') + return sampler + debug(f'Find sampler: name="{name}" found=None') + return None + + def list_samplers(): global all_samplers # pylint: disable=global-statement global all_samplers_map # pylint: disable=global-statement diff --git a/modules/ui_common.py b/modules/ui_common.py index d0696461a..84cae441d 100644 --- a/modules/ui_common.py +++ b/modules/ui_common.py @@ -54,7 +54,7 @@ def infotext_to_html(text): return code -def delete_files(js_data, files, _html_info, index): +def delete_files(js_data, files, all_files, index): try: data = json.loads(js_data) except Exception: @@ -63,25 +63,26 @@ def delete_files(js_data, files, _html_info, index): if index > -1 and shared.opts.save_selected_only and (index >= data['index_of_first_image']): files = [files[index]] start_index = index - filenames = [] - filenames = [] - fullfns = [] + deleted = [] + all_files = [f.split('/file=')[1] if 'file=' in f else f for f in all_files] if isinstance(all_files, list) else [] for _image_index, filedata in enumerate(files, start_index): - if 'name' in filedata and os.path.isfile(filedata['name']): - fullfn = filedata['name'] - filenames.append(os.path.basename(fullfn)) - try: - os.remove(fullfn) - base, _ext = os.path.splitext(fullfn) - desc = f'{base}.txt' - if os.path.exists(desc): - os.remove(desc) - fullfns.append(fullfn) - shared.log.info(f"Deleting image: {fullfn}") - except Exception as e: - shared.log.error(f'Error deleting file: {fullfn} {e}') - files = [image for image in files if image['name'] not in fullfns] - return files, plaintext_to_html(f"Deleted: {filenames[0] if len(filenames) > 0 else 'none'}") + try: + fn = filedata['name'] + if os.path.isfile(fn): + deleted.append(fn) + os.remove(fn) + if fn in all_files: + all_files.remove(fn) + shared.log.info(f'Delete: image="{fn}"') + base, _ext = os.path.splitext(fn) + desc = f'{base}.txt' + if os.path.exists(desc): + os.remove(desc) + shared.log.info(f'Delete: text="{fn}"') + except Exception as e: + shared.log.error(f'Delete: image="{fn}" {e}') + deleted = ', '.join(deleted) if len(deleted) > 0 else 'none' + return all_files, plaintext_to_html(f"Deleted: {deleted}") def save_files(js_data, files, html_info, index): @@ -296,8 +297,8 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe inputs=[generation_info, result_gallery, html_info, html_info], outputs=[download_files, html_log], ) - delete.click(fn=call_queue.wrap_gradio_call(delete_files),show_progress=False, - _js="(x, y, z, i) => [x, y, z, selected_gallery_index()]", + delete.click(fn=call_queue.wrap_gradio_call(delete_files), show_progress=False, + _js="(x, y, i, j) => [x, y, ...selected_gallery_files()]", inputs=[generation_info, result_gallery, html_info, html_info], outputs=[result_gallery, html_log], ) From 3d52e3fe9f6533b44ec2626f99011d61faa5d2eb Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 10:54:02 -0400 Subject: [PATCH 15/25] add /sdapi/v1/lora endpoint Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 +++ modules/api/api.py | 7 ++++--- modules/api/endpoints.py | 10 ---------- modules/api/loras.py | 22 ++++++++++++++++++++++ modules/lora/network.py | 21 +++++++++++++++++++++ 5 files changed, 50 insertions(+), 13 deletions(-) create mode 100644 modules/api/loras.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 34d357111..a7f02d718 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,9 @@ - Fix Dora loading - Remove per layer GC +- **API** + - Add `/sdapi/v1/lora?lora=` endpoint that returns full lora info and metadata + - **Fixes** - IPEX with DPM2++ FlowMatch samplers - Invalid attention processor with ControlNet diff --git a/modules/api/api.py b/modules/api/api.py index 72a2090a0..06b236d02 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -5,7 +5,7 @@ from fastapi import FastAPI, APIRouter, Depends, Request from fastapi.security import HTTPBasic, HTTPBasicCredentials from fastapi.exceptions import HTTPException from modules import errors, shared, postprocessing -from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, docs +from modules.api import models, endpoints, script, helpers, server, nvml, generate, process, control, gallery, loras, docs errors.install() @@ -100,8 +100,9 @@ class Api: # lora api if shared.native: - self.add_api_route("/sdapi/v1/loras", endpoints.get_loras, methods=["GET"], response_model=List[dict]) - self.add_api_route("/sdapi/v1/refresh-loras", endpoints.post_refresh_loras, methods=["POST"]) + self.add_api_route("/sdapi/v1/lora", loras.get_lora, methods=["GET"], response_model=dict) + self.add_api_route("/sdapi/v1/loras", loras.get_loras, methods=["GET"], response_model=List[dict]) + self.add_api_route("/sdapi/v1/refresh-loras", loras.post_refresh_loras, methods=["POST"]) # gallery api gallery.register_api(self.app) diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index 80b46f324..9dc94f92e 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -43,12 +43,6 @@ def get_embeddings(): return {"loaded": convert_embeddings(db.word_embeddings), "skipped": convert_embeddings(db.skipped_embeddings)} -def get_loras(): - from modules.lora import network, lora_load - def create_lora_json(obj: network.NetworkOnDisk): - return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata } - return [create_lora_json(obj) for obj in lora_load.available_networks.values()] - def get_extra_networks(page: Optional[str] = None, name: Optional[str] = None, filename: Optional[str] = None, title: Optional[str] = None, fullname: Optional[str] = None, hash: Optional[str] = None): # pylint: disable=redefined-builtin res = [] for pg in shared.extra_networks: @@ -158,10 +152,6 @@ def post_refresh_vae(): shared.refresh_vaes() return {} -def post_refresh_loras(): - from modules.lora import lora_load - return lora_load.list_available_networks() - def get_extensions_list(): from modules import extensions extensions.list_extensions() diff --git a/modules/api/loras.py b/modules/api/loras.py new file mode 100644 index 000000000..4dc7221d4 --- /dev/null +++ b/modules/api/loras.py @@ -0,0 +1,22 @@ +from fastapi.exceptions import HTTPException + + +def get_lora(lora: str) -> dict: + from modules.lora import lora_load + if lora not in lora_load.available_networks: + raise HTTPException(status_code=404, detail=f"Lora '{lora}' not found") + obj = lora_load.available_networks[lora] + obj.info = obj.get_info() + obj.desc = obj.get_desc() + print('HERE', obj) + return obj.__dict__ + +def get_loras(): + from modules.lora import network, lora_load + def create_lora_json(obj: network.NetworkOnDisk): + return { "name": obj.name, "alias": obj.alias, "path": obj.filename, "metadata": obj.metadata } + return [create_lora_json(obj) for obj in lora_load.available_networks.values()] + +def post_refresh_loras(): + from modules.lora import lora_load + return lora_load.list_available_networks() diff --git a/modules/lora/network.py b/modules/lora/network.py index 44e5a64b9..c22ae7cf9 100644 --- a/modules/lora/network.py +++ b/modules/lora/network.py @@ -90,6 +90,27 @@ class NetworkOnDisk: if not self.hash: self.set_hash(hashes.sha256(self.filename, "lora/" + self.name, use_addnet_hash=self.is_safetensors) or '') + def get_info(self): + data = {} + if shared.cmd_opts.no_metadata: + return data + if self.filename is not None: + fn = os.path.splitext(self.filename)[0] + '.json' + if os.path.exists(fn): + data = shared.readfile(fn, silent=True) + if type(data) is list: + data = data[0] + return data + + def get_desc(self): + if shared.cmd_opts.no_metadata: + return None + if self.filename is not None: + fn = os.path.splitext(self.filename)[0] + '.txt' + if os.path.exists(fn): + return shared.readfile(fn, silent=True) + return None + def get_alias(self): if shared.opts.lora_preferred_name == "filename": return self.name From 6c3f0dd43bcb6dfa9aac8fcea74ac92987a1ab29 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 11:23:39 -0400 Subject: [PATCH 16/25] add `SD_SAVE_DEBUG` env variable Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + javascript/gallery.js | 2 +- modules/images.py | 6 ++++++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a7f02d718..cd9c1a77b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,6 +26,7 @@ - Add `wheel` to requirements due to `pip` change - Case-insensitive sampler name matching - Fix delete file with gallery views + - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen ## Update for 2025-06-16 diff --git a/javascript/gallery.js b/javascript/gallery.js index 54945cf97..eaa646c66 100644 --- a/javascript/gallery.js +++ b/javascript/gallery.js @@ -131,7 +131,7 @@ class GalleryFile extends HTMLElement { } const ext = this.name.split('.').pop().toLowerCase(); if (!['jpg', 'jpeg', 'png', 'gif', 'webp', 'jxl', 'svg', 'mp4'].includes(ext)) { - console.error(`gallery: type=${ext} file=${this.name} unsupported`); + // console.error(`gallery: type=${ext} file=${this.name} unsupported`); return; } this.hash = await getHash(`${this.folder}/${this.name}/${this.size}/${this.mtime}`); // eslint-disable-line no-use-before-define diff --git a/modules/images.py b/modules/images.py index e8a0fa384..d5c2d3130 100644 --- a/modules/images.py +++ b/modules/images.py @@ -19,6 +19,7 @@ from modules.video import save_video # pylint: disable=unused-import debug = errors.log.trace if os.environ.get('SD_PATH_DEBUG', None) is not None else lambda *args, **kwargs: None +debug_save = errors.log.trace if os.environ.get('SD_SAVE_DEBUG', None) is not None else lambda *args, **kwargs: None try: from pi_heif import register_heif_opener register_heif_opener() @@ -73,6 +74,7 @@ def atomically_save_image(): pnginfo_data = PngImagePlugin.PngInfo() for k, v in params.pnginfo.items(): pnginfo_data.add_text(k, str(v)) + debug_save(f'Save pnginfo: {params.pnginfo.items()}') save_args = { 'compress_level': 6, 'pnginfo': pnginfo_data if shared.opts.image_metadata else None } elif image_format == 'JPEG': if image.mode == 'RGBA': @@ -82,12 +84,14 @@ def atomically_save_image(): image = image.point(lambda p: p * 0.0038910505836576).convert("L") save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality } if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) elif image_format == 'WEBP': if image.mode == 'I;16': image = image.point(lambda p: p * 0.0038910505836576).convert("RGB") save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) elif image_format == 'JXL': if image.mode == 'I;16': @@ -96,10 +100,12 @@ def atomically_save_image(): image = image.convert("RGBA") save_args = { 'optimize': True, 'quality': shared.opts.jpeg_quality, 'lossless': shared.opts.webp_lossless } if shared.opts.image_metadata: + debug_save(f'Save exif: {exifinfo}') save_args['exif'] = piexif.dump({ "Exif": { piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(exifinfo, encoding="unicode") } }) else: save_args = { 'quality': shared.opts.jpeg_quality } try: + debug_save(f'Save args: {save_args}') image.save(fn, format=image_format, **save_args) except Exception as e: shared.log.error(f'Save failed: file="{fn}" format={image_format} args={save_args} {e}') From 5b486a6ef17d0db34ac2f8f7c922adbd6807555f Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 15:32:37 -0400 Subject: [PATCH 17/25] sdnq add xyz grid support, improve offloading compatibility Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 7 +++++- modules/api/loras.py | 1 - modules/hidiffusion/__init__.py | 2 +- modules/interrogate/joycaption.py | 5 +++- modules/interrogate/vqa.py | 5 ++-- modules/ipadapter.py | 2 +- modules/model_quant.py | 19 ++++++++++---- modules/model_te.py | 16 +++++++----- modules/sd_models.py | 5 ++-- modules/sd_offload.py | 21 +++++++++++++++- modules/sd_vae_taesd.py | 42 +++++++++++++++++-------------- modules/shared.py | 5 ++-- modules/timer.py | 5 ++-- scripts/xyz_grid_classes.py | 10 +++++++- scripts/xyz_grid_shared.py | 12 +++++++++ 15 files changed, 112 insertions(+), 45 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cd9c1a77b..a4259c5bd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,14 +3,18 @@ ## Update for 2025-06-25 - **Changes** - - Use Diffusers version of *OmniGen* + - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support - Support Remote VAE with *Omnigen, Lumina 2 and PixArt* + - Use Diffusers version of *OmniGen* - **SDNQ Quantization** - Add modules_to_not_convert support for post mode - Fix Qwen 2.5 with int8 matmul - Fix Dora loading - Remove per layer GC + - Improve offload compatibility + - Add support for XYZ grid to test quantization modes + *note*: you need to enable quantization and choose what it applies on, then xyz grid can change quantization mode - **API** - Add `/sdapi/v1/lora?lora=` endpoint that returns full lora info and metadata @@ -27,6 +31,7 @@ - Case-insensitive sampler name matching - Fix delete file with gallery views - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen + - Fix TAESD model type detection ## Update for 2025-06-16 diff --git a/modules/api/loras.py b/modules/api/loras.py index 4dc7221d4..7e65a709d 100644 --- a/modules/api/loras.py +++ b/modules/api/loras.py @@ -8,7 +8,6 @@ def get_lora(lora: str) -> dict: obj = lora_load.available_networks[lora] obj.info = obj.get_info() obj.desc = obj.get_desc() - print('HERE', obj) return obj.__dict__ def get_loras(): diff --git a/modules/hidiffusion/__init__.py b/modules/hidiffusion/__init__.py index ac7dd9627..2e1a500d9 100644 --- a/modules/hidiffusion/__init__.py +++ b/modules/hidiffusion/__init__.py @@ -41,5 +41,5 @@ def apply(p, model_type): def unapply(): pipe = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model - if hasattr(pipe, 'unet'): + if hasattr(pipe, 'unet') and pipe.unet is not None: hidiffusion.remove_hidiffusion(pipe) diff --git a/modules/interrogate/joycaption.py b/modules/interrogate/joycaption.py index cc316341a..4941f7899 100644 --- a/modules/interrogate/joycaption.py +++ b/modules/interrogate/joycaption.py @@ -58,9 +58,12 @@ opts = JoyOptions() @torch.no_grad() -def predict(question: str, image): +def predict(question: str, image, vqa_model: str = None) -> str: global llava_model, processor # pylint: disable=global-statement opts.max_new_tokens = shared.opts.interrogate_vlm_max_length + if vqa_model is not None and opts.repo != vqa_model: + opts.repo = vqa_model + llava_model = None if llava_model is None: shared.log.info(f'Interrogate: type=vlm model="JoyCaption" {str(opts)}') processor = AutoProcessor.from_pretrained(opts.repo) diff --git a/modules/interrogate/vqa.py b/modules/interrogate/vqa.py index 32820a243..5d74a8b9c 100644 --- a/modules/interrogate/vqa.py +++ b/modules/interrogate/vqa.py @@ -38,7 +38,8 @@ vlm_models = { "ToriiGate 0.4 2B": "Minthy/ToriiGate-v0.4-2B", "ToriiGate 0.4 7B": "Minthy/ToriiGate-v0.4-7B", "ViLT Base": "dandelin/vilt-b32-finetuned-vqa", # 0.5GB - "JoyCaption": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB + "JoyCaption Alpha": "fancyfeast/llama-joycaption-alpha-two-hf-llava", # 17.4GB + "JoyCaption Beta": "fancyfeast/llama-joycaption-beta-one-hf-llava", # 17.4GB "JoyTag": "fancyfeast/joytag", # 0.7GB "AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B", "AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B", @@ -583,7 +584,7 @@ def interrogate(question:str='', system_prompt:str=None, prompt:str=None, image: answer = joytag.predict(image) elif 'joycaption' in vqa_model.lower(): from modules.interrogate import joycaption - answer = joycaption.predict(question, image) + answer = joycaption.predict(question, image, vqa_model) elif 'deepseek' in vqa_model.lower(): from modules.interrogate import deepseek answer = deepseek.predict(question, image, vqa_model) diff --git a/modules/ipadapter.py b/modules/ipadapter.py index a03381a1b..3b2ce3ea9 100644 --- a/modules/ipadapter.py +++ b/modules/ipadapter.py @@ -146,7 +146,7 @@ def unapply(pipe, unload: bool = False): # pylint: disable=arguments-differ if unload: shared.log.debug('IP adapter unload') pipe.unload_ip_adapter() - if hasattr(pipe, 'unet'): + if hasattr(pipe, 'unet') and pipe.unet is not None: module = pipe.unet elif hasattr(pipe, 'transformer'): module = pipe.transformer diff --git a/modules/model_quant.py b/modules/model_quant.py index b3db3f311..1eab2aa64 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -113,11 +113,15 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig - if weights_dtype is None: - if shared.opts.sdnq_quantize_weights_mode_te != "default" and module in {"TE", "LLM"}: - weights_dtype = shared.opts.sdnq_quantize_weights_mode_te - else: + if weights_dtype is None and module in {"TE", "LLM"}: + if shared.opts.sdnq_quantize_weights_mode_te == "none": + return None + elif shared.opts.sdnq_quantize_weights_mode_te == "same as model" or shared.opts.sdnq_quantize_weights_mode_te == "default": weights_dtype = shared.opts.sdnq_quantize_weights_mode + else: + weights_dtype = shared.opts.sdnq_quantize_weights_mode_te + if weights_dtype is None: + return None if shared.opts.device_map == "gpu": quantization_device = devices.device @@ -337,6 +341,11 @@ def sdnq_quantize_model(model, op=None, sd_model=None, do_gc=True): else: weights_dtype = shared.opts.sdnq_quantize_weights_mode + if weights_dtype is None or weights_dtype == 'none': + return model + if debug: + log.trace(f'Quantization: type=SDNQ op={op} cls={model.__class__} dtype={weights_dtype} mode{shared.opts.diffusers_offload_mode}') + if shared.opts.diffusers_offload_mode in {"none", "model"}: quantization_device = devices.device if shared.opts.sdnq_quantize_with_gpu else devices.cpu return_device = devices.device @@ -397,7 +406,7 @@ def sdnq_quantize_weights(sd_model): try: t0 = time.time() from modules import shared, devices, sd_models - log.info(f"Quantization: type=SDNQ dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32} modules={shared.opts.sdnq_quantize_weights}") + log.debug(f"Quantization: type=SDNQ modules={shared.opts.sdnq_quantize_weights} dtype={shared.opts.sdnq_quantize_weights_mode} dtype_te={shared.opts.sdnq_quantize_weights_mode_te} matmul={shared.opts.sdnq_use_quantized_matmul} group_size={shared.opts.sdnq_quantize_weights_group_size} quant_conv={shared.opts.sdnq_quantize_conv_layers} matmul_conv={shared.opts.sdnq_use_quantized_matmul_conv} quantize_with_gpu={shared.opts.sdnq_quantize_with_gpu} dequantize_fp32={shared.opts.sdnq_dequantize_fp32}") global quant_last_model_name, quant_last_model_device # pylint: disable=global-statement sd_model = sd_models.apply_function_to_model(sd_model, sdnq_quantize_model, shared.opts.sdnq_quantize_weights, op="sdnq") diff --git a/modules/model_te.py b/modules/model_te.py index 6472721ee..27ad10357 100644 --- a/modules/model_te.py +++ b/modules/model_te.py @@ -4,7 +4,6 @@ import torch import transformers from safetensors.torch import load_file from modules import shared, devices, files_cache, errors, model_quant -from installer import install te_dict = {} @@ -72,27 +71,32 @@ def load_t5(name=None, cache_dir=None): elif 'int8' in name.lower(): from modules.model_quant import create_sdnq_config quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='int8') - t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) + if quantization_config is not None: + t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif 'uint4' in name.lower(): from modules.model_quant import create_sdnq_config quantization_config = create_sdnq_config(kwargs=None, allow_sdnq=True, module='any', weights_dtype='uint4') - t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) + if quantization_config is not None: + t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif 'qint4' in name.lower(): model_quant.load_quanto('Load model: type=T5') quantization_config = transformers.QuantoConfig(weights='int4') - t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) + if quantization_config is not None: + t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif 'qint8' in name.lower(): model_quant.load_quanto('Load model: type=T5') quantization_config = transformers.QuantoConfig(weights='int8') - t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) + if quantization_config is not None: + t5 = transformers.T5EncoderModel.from_pretrained(repo_id, subfolder='text_encoder_3', quantization_config=quantization_config, cache_dir=cache_dir, torch_dtype=devices.dtype) elif '/' in name: shared.log.debug(f'Load model: type=T5 repo={name}') quant_config = model_quant.create_config(module='TE') - t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config) + if quantization_config is not None: + t5 = transformers.T5EncoderModel.from_pretrained(name, cache_dir=cache_dir, torch_dtype=devices.dtype, **quant_config) else: t5 = None diff --git a/modules/sd_models.py b/modules/sd_models.py index 88644b17e..f6a7e53db 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -660,7 +660,7 @@ def load_diffuser(checkpoint_info=None, already_loaded_state_dict=None, timer=No from modules import modelstats modelstats.analyze() - shared.log.info(f"Load {op}: time={timer.summary()} native={get_native(sd_model)} memory={memory_stats()}") + shared.log.info(f"Load {op}: family={shared.sd_model_type} time={timer.dct()} native={get_native(sd_model)} memory={memory_stats()}") class DiffusersTaskType(Enum): @@ -1086,7 +1086,8 @@ def clear_caches(): lora_common.loaded_networks.clear() lora_common.previously_loaded_networks.clear() lora_load.lora_cache.clear() - from modules import prompt_parser_diffusers, memstats + from modules import prompt_parser_diffusers, memstats, sd_offload + sd_offload.offload_hook_instance = None prompt_parser_diffusers.cache.clear() memstats.reset_stats() diff --git a/modules/sd_offload.py b/modules/sd_offload.py index 57057d233..e36db4332 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -4,6 +4,7 @@ import time import inspect import torch import accelerate.hooks +import accelerate.utils.modeling from installer import log from modules import shared, devices, errors, model_quant from modules.timer import process as process_timer @@ -15,6 +16,16 @@ offload_warn = ['sc', 'sd3', 'f1', 'h1', 'hunyuandit', 'auraflow', 'omnigen', 'c offload_post = ['h1'] offload_hook_instance = None balanced_offload_exclude = ['CogView4Pipeline'] +accelerate_dtype_byte_size = None + + +def dtype_byte_size(dtype: torch.dtype): + try: + if dtype in [torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz]: + dtype = accelerate.utils.modeling.CustomDtype.FP8 + except Exception: # catch since older torch many not have defined dtypes + pass + return accelerate_dtype_byte_size(dtype) def get_signature(cls): @@ -58,6 +69,7 @@ def set_accelerate(sd_model): def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False): + global accelerate_dtype_byte_size # pylint: disable=global-statement t0 = time.time() if not shared.native: shared.log.warning('Attempting to use offload with backend=original') @@ -67,6 +79,9 @@ def set_diffuser_offload(sd_model, op:str='model', quiet:bool=False): return if not (hasattr(sd_model, "has_accelerate") and sd_model.has_accelerate): sd_model.has_accelerate = False + if accelerate_dtype_byte_size is None: + accelerate_dtype_byte_size = accelerate.utils.modeling.dtype_byte_size + accelerate.utils.modeling.dtype_byte_size = dtype_byte_size if shared.opts.diffusers_offload_mode == "none": if shared.sd_model_type in offload_warn or 'video' in shared.sd_model_type: shared.log.warning(f'Setting {op}: offload={shared.opts.diffusers_offload_mode} type={shared.sd_model.__class__.__name__} large model') @@ -163,14 +178,18 @@ class OffloadHook(accelerate.hooks.ModelHook): max_memory = { device_index: self.gpu, "cpu": self.cpu } device_map = getattr(module, "balanced_offload_device_map", None) if device_map is None or max_memory != getattr(module, "balanced_offload_max_memory", None): + # try: device_map = accelerate.infer_auto_device_map(module, max_memory=max_memory) + # except Exception as e: + # shared.log.error(f'Offload: type=balanced module={module.__class__.__name__} {e}') offload_dir = getattr(module, "offload_dir", os.path.join(shared.opts.accelerate_offload_path, module.__class__.__name__)) if devices.backend == "directml": keys = device_map.keys() for v in keys: if isinstance(device_map[v], int): device_map[v] = f"{devices.device.type}:{device_map[v]}" # int implies CUDA or XPU device, but it will break DirectML backend so we add type - module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir) + if device_map is not None: + module = accelerate.dispatch_model(module, device_map=device_map, offload_dir=offload_dir) module._hf_hook.execution_device = torch.device(devices.device) # pylint: disable=protected-access module.balanced_offload_device_map = device_map module.balanced_offload_max_memory = max_memory diff --git a/modules/sd_vae_taesd.py b/modules/sd_vae_taesd.py index 7837c2b14..08fd834f6 100644 --- a/modules/sd_vae_taesd.py +++ b/modules/sd_vae_taesd.py @@ -52,34 +52,36 @@ def warn_once(msg, variant=None): def get_model(model_type = 'decoder', variant = None): global prev_cls, prev_type, prev_model # pylint: disable=global-statement from modules import shared - cls = shared.sd_model_type - if cls in {'ldm', 'pixartalpha'}: - cls = 'sd' - elif cls in {'h1', 'lumina2'}: - cls = 'f1' - elif cls in {'pixartsigma', 'omnigen'}: - cls = 'sdxl' - elif cls not in supported: - warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant) + model_cls = shared.sd_model_type + if model_cls is None or model_cls == 'none': + return None + elif model_cls in {'ldm', 'pixartalpha'}: + model_cls = 'sd' + elif model_cls in {'h1', 'lumina2'}: + model_cls = 'f1' + elif model_cls in {'pixartsigma', 'omnigen'}: + model_cls = 'sdxl' + elif model_cls not in supported: + warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant) variant = variant or shared.opts.taesd_variant folder = os.path.join(paths.models_path, "TAESD") os.makedirs(folder, exist_ok=True) if variant.startswith('TAE'): cfg = TAESD_MODELS[variant] - if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): + if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] - fn = os.path.join(folder, cfg['fn'] + cls + '_' + model_type + '.pth') + fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth') if not os.path.exists(fn): uri = cfg['uri'] if not uri.endswith('.pth'): - uri += '/tae' + cls + '_' + model_type + '.pth' + uri += '/tae' + model_cls + '_' + model_type + '.pth' try: shared.log.info(f'Decode: type="taesd" variant="{variant}": uri="{uri}" fn="{fn}" download') torch.hub.download_url_to_file(uri, fn) except Exception as e: warn_once(f'download uri={uri} {e}', variant=variant) if os.path.exists(fn): - prev_cls = cls + prev_cls = model_cls prev_type = model_type prev_model = variant shared.log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" load') @@ -97,14 +99,14 @@ def get_model(model_type = 'decoder', variant = None): TAESD_MODELS[variant]['model'] = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None) return TAESD_MODELS[variant]['model'] elif variant.startswith('Hybrid'): - cfg = CQYAN_MODELS[variant].get(cls, None) - if (cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): + cfg = CQYAN_MODELS[variant].get(model_cls, None) + if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None): return cfg['model'] if cfg is None: - warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant) + warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant) return None repo = cfg['repo'] - prev_cls = cls + prev_cls = model_cls prev_type = model_type prev_model = variant shared.log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load') @@ -116,10 +118,12 @@ def get_model(model_type = 'decoder', variant = None): from modules.taesd.hybrid_small import AutoencoderSmall vae = AutoencoderSmall.from_pretrained(repo, cache_dir=shared.opts.hfcache_dir, torch_dtype=dtype) vae = vae.to(devices.device, dtype=dtype) - CQYAN_MODELS[variant][cls]['model'] = vae + CQYAN_MODELS[variant][model_cls]['model'] = vae return vae + elif variant is None: + warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} variant is none', variant=variant) else: - warn_once(f'cls={shared.sd_model.__class__.__name__} type={cls} unsuppported', variant=variant) + warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant) return None diff --git a/modules/shared.py b/modules/shared.py index 2231e0484..276143e42 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -66,6 +66,7 @@ dir_timestamps = {} dir_cache = {} max_workers = 8 default_hfcache_dir = os.environ.get("SD_HFCACHEDIR", None) or os.path.join(os.path.expanduser('~'), '.cache', 'huggingface', 'hub') +sdnq_quant_modes = ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"] class Backend(Enum): @@ -518,8 +519,8 @@ options_templates.update(options_section(("quantization", "Quantization Settings "sdnq_quantize_sep": OptionInfo("

SDNQ: SD.Next Quantization

", "", gr.HTML), "sdnq_quantize_weights": OptionInfo([], "Quantization enabled", gr.CheckboxGroup, {"choices": ["Model", "Transformer", "VAE", "TE", "Video", "LLM", "ControlNet"], "visible": native}), "sdnq_quantize_mode": OptionInfo("pre", "Quantization mode", gr.Dropdown, {"choices": ["pre", "post"], "visible": native}), - "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": ["int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}), - "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ["default", "int8", "float8_e4m3fn", "int7", "int6", "int5", "uint4", "uint3", "uint2", "float8_e5m2", "float8_e4m3fnuz", "float8_e5m2fnuz", "uint8", "uint7", "uint6", "uint5", "int4", "int3", "int2", "uint1"], "visible": native}), + "sdnq_quantize_weights_mode": OptionInfo("int8", "Quantization type", gr.Dropdown, {"choices": sdnq_quant_modes, "visible": native}), + "sdnq_quantize_weights_mode_te": OptionInfo("default", "Quantization type for Text Encoders", gr.Dropdown, {"choices": ['default'] + sdnq_quant_modes, "visible": native}), "sdnq_quantize_weights_group_size": OptionInfo(0, "Group size", gr.Slider, {"minimum": -1, "maximum": 4096, "step": 1, "visible": native}), "sdnq_quantize_conv_layers": OptionInfo(False, "Quantize the convolutional layers", gr.Checkbox, {"visible": native}), "sdnq_dequantize_compile": OptionInfo(devices.has_triton(), "Dequantize using torch.compile", gr.Checkbox, {"visible": native}), diff --git a/modules/timer.py b/modules/timer.py index 69107e605..59c6a1de3 100644 --- a/modules/timer.py +++ b/modules/timer.py @@ -44,8 +44,7 @@ class Timer: def summary(self, min_time=default_min_time, total=True): if self.profile: min_time = -1 - if self.total <= 0: - self.total = sum(self.records.values()) + self.total = sum(self.records.values()) res = f"total={self.total:.2f} " if total else '' additions = [x for x in self.records.items() if x[1] >= min_time] additions = sorted(additions, key=lambda x: x[1], reverse=True) @@ -60,6 +59,8 @@ class Timer: def dct(self, min_time=default_min_time): if self.profile: res = {k: round(v, 4) for k, v in self.records.items()} + self.total = sum(self.records.values()) + self.records['total'] = self.total res = {k: round(v, 2) for k, v in self.records.items() if v >= min_time} res = {k: v for k, v in sorted(res.items(), key=lambda x: x[1], reverse=True)} # noqa: C416 # pylint: disable=unnecessary-comprehension return res diff --git a/scripts/xyz_grid_classes.py b/scripts/xyz_grid_classes.py index 0d3d6dd97..289059afb 100644 --- a/scripts/xyz_grid_classes.py +++ b/scripts/xyz_grid_classes.py @@ -1,4 +1,4 @@ -from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import +from scripts.xyz_grid_shared import apply_field, apply_task_arg, apply_task_args, apply_setting, apply_prompt_primary, apply_prompt_refine, apply_prompt_detailer, apply_prompt_all, apply_order, apply_sampler, apply_hr_sampler_name, confirm_samplers, apply_checkpoint, apply_refiner, apply_unet, apply_dict, apply_clip_skip, apply_vae, list_lora, apply_lora, apply_lora_strength, apply_te, apply_styles, apply_upscaler, apply_context, apply_detailer, apply_override, apply_processing, apply_options, apply_seed, apply_sdnq_quant, apply_sdnq_quant_te, format_value_add_label, format_bool, format_value, format_value_join_list, do_nothing, format_nothing, str_permutations # pylint: disable=no-name-in-module, unused-import from modules import shared, shared_items, sd_samplers, ipadapter, sd_models, sd_vae, sd_unet @@ -58,6 +58,7 @@ class SharedSettingsStackHelper(object): extra_networks_default_multiplier = None disable_apply_metadata = None disable_apply_params = None + sdnq_quant_mode = None def __enter__(self): # Save overridden settings so they can be restored later @@ -89,6 +90,8 @@ class SharedSettingsStackHelper(object): self.teacache_thresh = shared.opts.teacache_thresh self.disable_apply_metadata = shared.opts.disable_apply_metadata self.disable_apply_params = shared.opts.disable_apply_params + self.sdnq_quant_mode = shared.opts.sdnq_quantize_weights_mode + shared.opts.data["disable_apply_metadata"] = [] shared.opts.data["disable_apply_params"] = '' @@ -135,6 +138,9 @@ class SharedSettingsStackHelper(object): if self.sd_unet != shared.opts.sd_unet: shared.opts.data["sd_unet"] = self.sd_unet sd_unet.load_unet(shared.sd_model) + if self.sdnq_quant_mode != shared.opts.sdnq_quantize_weights_mode: + shared.opts.data["sdnq_quantize_weights_mode"] = self.sdnq_quant_mode + sd_models.reload_model_weights(op='model') axis_options = [ @@ -193,6 +199,8 @@ axis_options = [ AxisOption("[Postprocess] Context", str, apply_context, choices=lambda: ["Add with forward", "Remove with forward", "Add with backward", "Remove with backward"]), AxisOption("[Postprocess] Detailer", str, apply_detailer, fmt=format_value_add_label), AxisOption("[Postprocess] Detailer strength", str, apply_field("detailer_strength")), + AxisOption("[Quant] SDNQ quant mode", str, apply_sdnq_quant, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)), + AxisOption("[Quant] SDNQ quant mode TE", str, apply_sdnq_quant_te, cost=0.9, fmt=format_value_add_label, choices=lambda: ['none'] + sorted(shared.sdnq_quant_modes)), AxisOption("[HDR] Mode", int, apply_field("hdr_mode")), AxisOption("[HDR] Brightness", float, apply_field("hdr_brightness")), AxisOption("[HDR] Color", float, apply_field("hdr_color")), diff --git a/scripts/xyz_grid_shared.py b/scripts/xyz_grid_shared.py index 2d594bac7..265af6862 100644 --- a/scripts/xyz_grid_shared.py +++ b/scripts/xyz_grid_shared.py @@ -147,6 +147,18 @@ def confirm_samplers(p, xs): shared.log.warning(f"XYZ grid: unknown sampler: {x}") +def apply_sdnq_quant(p, x, xs): + shared.opts.sdnq_quantize_weights_mode = x + sd_models.unload_model_weights(op='model') # reload will happen on-demand + shared.log.debug(f'XYZ grid apply sdnq quant: mode="{x}"') + + +def apply_sdnq_quant_te(p, x, xs): + shared.opts.sdnq_quantize_weights_mode_te = x + sd_models.unload_model_weights(op='model') # reload will happen on-demand + shared.log.debug(f'XYZ grid apply sdnq quant te: mode="{x}"') + + def apply_checkpoint(p, x, xs): if x == shared.opts.sd_model_checkpoint: return From 3a1d98c472041e5274d1d19931388f3a26f05911 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 15:33:20 -0400 Subject: [PATCH 18/25] add joycaption beta support Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4259c5bd..d4761d4ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,7 +3,7 @@ ## Update for 2025-06-25 - **Changes** - - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support + - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support (in addition to existing JoyCaption Alpha) - Support Remote VAE with *Omnigen, Lumina 2 and PixArt* - Use Diffusers version of *OmniGen* From f8977d2f01b2c34e0d18f75f1e728d7970019307 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 15:54:43 -0400 Subject: [PATCH 19/25] add /sdapi/v1/controlnets api endpoint Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + modules/api/api.py | 1 + modules/api/endpoints.py | 4 ++++ modules/control/units/controlnet.py | 17 +++++++++++++++++ 4 files changed, 23 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d4761d4ac..a1ed3e66a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +18,7 @@ - **API** - Add `/sdapi/v1/lora?lora=` endpoint that returns full lora info and metadata + - Add `/sdapi/v1/controlnets?model_type=` endpoints that returns list of available controlnets for specific model type - **Fixes** - IPEX with DPM2++ FlowMatch samplers diff --git a/modules/api/api.py b/modules/api/api.py index 06b236d02..ad26a384a 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -78,6 +78,7 @@ class Api: self.add_api_route("/sdapi/v1/samplers", endpoints.get_samplers, methods=["GET"], response_model=List[models.ItemSampler]) self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler]) self.add_api_route("/sdapi/v1/sd-models", endpoints.get_sd_models, methods=["GET"], response_model=List[models.ItemModel]) + self.add_api_route("/sdapi/v1/controlnets", endpoints.get_controlnets, methods=["GET"], response_model=List[str]) self.add_api_route("/sdapi/v1/hypernetworks", endpoints.get_hypernetworks, methods=["GET"], response_model=List[models.ItemHypernetwork]) self.add_api_route("/sdapi/v1/face-restorers", endpoints.get_detailers, methods=["GET"], response_model=List[models.ItemDetailer]) self.add_api_route("/sdapi/v1/prompt-styles", endpoints.get_prompt_styles, methods=["GET"], response_model=List[models.ItemStyle]) diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index 9dc94f92e..87181ba4f 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -23,6 +23,10 @@ def get_sd_models(): checkpoints.append({"title": v.title, "model_name": v.name, "filename": v.filename, "type": v.type, "hash": v.shorthash, "sha256": v.sha256, "config": sd_models_config.find_checkpoint_config_near_filename(v)}) return checkpoints +def get_controlnets(model_type: Optional[str] = None): + from modules.control.units.controlnet import api_list_models + return api_list_models() + def get_hypernetworks(): return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks] diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index d2e59348c..debe54c16 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -137,6 +137,23 @@ def find_models(): find_models() + +def api_list_models(model_type: str = None): + import modules.shared + model_type = model_type or modules.shared.sd_model_type + model_list = [] + if model_type == 'sd' or model_type == 'all': + model_list += list(predefined_sd15) + if model_type == 'sdxl' or model_type == 'all': + model_list += list(predefined_sdxl) + if model_type == 'f1' or model_type == 'all': + model_list += list(predefined_f1) + if model_type == 'sd3' or model_type == 'all': + model_list += list(predefined_sd3) + model_list += sorted(find_models()) + return model_list + + def list_models(refresh=False): import modules.shared global models # pylint: disable=global-statement From 81e55f04592e29583ad6ebab160060e9c7d66c5b Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 25 Jun 2025 22:58:40 +0300 Subject: [PATCH 20/25] Fix SDNQ pre-mode --- modules/model_quant.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/modules/model_quant.py b/modules/model_quant.py index 1eab2aa64..c8fd0f9b7 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -113,15 +113,18 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig - if weights_dtype is None and module in {"TE", "LLM"}: - if shared.opts.sdnq_quantize_weights_mode_te == "none": - return None - elif shared.opts.sdnq_quantize_weights_mode_te == "same as model" or shared.opts.sdnq_quantize_weights_mode_te == "default": - weights_dtype = shared.opts.sdnq_quantize_weights_mode - else: - weights_dtype = shared.opts.sdnq_quantize_weights_mode_te if weights_dtype is None: - return None + if module in {"TE", "LLM"}: + if shared.opts.sdnq_quantize_weights_mode_te == "none": + return kwargs + elif shared.opts.sdnq_quantize_weights_mode_te in {"same as model", "default"}: + weights_dtype = shared.opts.sdnq_quantize_weights_mode + else: + weights_dtype = shared.opts.sdnq_quantize_weights_mode_te + elif shared.opts.sdnq_quantize_weights_mode == "none": + return kwargs + else: + weights_dtype = shared.opts.sdnq_quantize_weights_mode if shared.opts.device_map == "gpu": quantization_device = devices.device From e43d1d2ba7308df035db2d5c7b64a8d2d2388bc4 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 25 Jun 2025 23:25:49 +0300 Subject: [PATCH 21/25] SDNQ use strings as target_dtype --- modules/sdnq/common.py | 31 ++++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py index ccfdf1227..ab8b70c0d 100644 --- a/modules/sdnq/common.py +++ b/modules/sdnq/common.py @@ -2,33 +2,34 @@ import sys import torch -from accelerate.utils import CustomDtype from modules import devices torch_version = float(torch.__version__[:3]) dtype_dict = { "int8": {"min": -128, "max": 127, "num_bits": 8, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.int8, "is_unsigned": False, "is_integer": True}, - "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": torch.int8, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, - "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int7": {"min": -64, "max": 63, "num_bits": 7, "target_dtype": "int7", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int6": {"min": -32, "max": 31, "num_bits": 6, "target_dtype": "int6", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int5": {"min": -16, "max": 15, "num_bits": 5, "target_dtype": "int5", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int4": {"min": -8, "max": 7, "num_bits": 4, "target_dtype": "int4", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int3": {"min": -4, "max": 3, "num_bits": 3, "target_dtype": "int3", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, + "int2": {"min": -2, "max": 1, "num_bits": 2, "target_dtype": "int2", "torch_dtype": torch.int8, "storage_dtype": torch.uint8, "is_unsigned": False, "is_integer": True}, "uint8": {"min": 0, "max": 255, "num_bits": 8, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": torch.uint8, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": CustomDtype.INT4, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, - "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": CustomDtype.INT2, "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint7": {"min": 0, "max": 127, "num_bits": 7, "target_dtype": "uint7", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint6": {"min": 0, "max": 63, "num_bits": 6, "target_dtype": "uint6", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint5": {"min": 0, "max": 31, "num_bits": 5, "target_dtype": "uint5", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint4": {"min": 0, "max": 15, "num_bits": 4, "target_dtype": "uint4", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint3": {"min": 0, "max": 7, "num_bits": 3, "target_dtype": "uint3", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, + "uint2": {"min": 0, "max": 3, "num_bits": 2, "target_dtype": "uint2", "torch_dtype": torch.uint8, "storage_dtype": torch.uint8, "is_unsigned": True, "is_integer": True}, "uint1": {"min": 0, "max": 1, "num_bits": 1, "target_dtype": torch.bool, "torch_dtype": torch.bool, "storage_dtype": torch.bool, "is_unsigned": True, "is_integer": True}, "float8_e4m3fn": {"min": -448, "max": 448, "num_bits": 8, "target_dtype": torch.float8_e4m3fn, "torch_dtype": torch.float8_e4m3fn, "storage_dtype": torch.float8_e4m3fn, "is_unsigned": False, "is_integer": False}, "float8_e5m2": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": torch.float8_e5m2, "torch_dtype": torch.float8_e5m2, "storage_dtype": torch.float8_e5m2, "is_unsigned": False, "is_integer": False}, - "float8_e4m3fnuz": {"min": -240, "max": 240, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False}, - "float8_e5m2fnuz": {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": CustomDtype.FP8, "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False}, } dtype_dict["bool"] = dtype_dict["uint1"] +if hasattr(torch, "float8_e4m3fnuz"): + dtype_dict["float8_e4m3fnuz"] = {"min": -240, "max": 240, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e4m3fnuz, "storage_dtype": torch.float8_e4m3fnuz, "is_unsigned": False, "is_integer": False} +if hasattr(torch, "float8_e5m2fnuz"): + dtype_dict["float8_e5m2fnuz"] = {"min": -57344, "max": 57344, "num_bits": 8, "target_dtype": "fp8", "torch_dtype": torch.float8_e5m2fnuz, "storage_dtype": torch.float8_e5m2fnuz, "is_unsigned": False, "is_integer": False} use_tensorwise_fp8_matmul = torch_version < 2.5 or devices.backend in {"cpu", "openvino"} or (devices.backend == "cuda" and sys.platform == "win32" and torch_version <= 2.7 and torch.cuda.get_device_capability(devices.device) == (8,9)) quantized_matmul_dtypes = ("int8", "int7", "int6", "int5", "int4", "int3", "int2", "float8_e4m3fn", "float8_e5m2") From c9f49720c55cd7bbe2fb620e33418a45e00892db Mon Sep 17 00:00:00 2001 From: Disty0 Date: Wed, 25 Jun 2025 23:37:41 +0300 Subject: [PATCH 22/25] Cleanup --- modules/devices.py | 2 +- modules/sd_hijack_accelerate.py | 4 ++-- modules/sd_offload.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/modules/devices.py b/modules/devices.py index 1c35f2683..8d15fd238 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -665,6 +665,6 @@ def normalize_device(dev): def same_device(d1, d2): - if d1.type != d2.type: + if torch.device(d1).type != torch.device(d2).type: return False return normalize_device(d1) == normalize_device(d2) diff --git a/modules/sd_hijack_accelerate.py b/modules/sd_hijack_accelerate.py index f8cf8983f..7f312a029 100644 --- a/modules/sd_hijack_accelerate.py +++ b/modules/sd_hijack_accelerate.py @@ -36,7 +36,7 @@ def hijack_set_module_tensor( # note: majority of time is spent on .to(old_value.dtype) if tensor_name in module._buffers: # pylint: disable=protected-access module._buffers[tensor_name] = value.to(device, old_value.dtype) # pylint: disable=protected-access - elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access + elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device, old_value.dtype) # pylint: disable=protected-access t1 = time.time() @@ -64,7 +64,7 @@ def hijack_set_module_tensor_simple( with devices.inference_context(): if tensor_name in module._buffers: # pylint: disable=protected-access module._buffers[tensor_name] = value.to(device) # pylint: disable=protected-access - elif value is not None or not devices.same_device(torch.device(device), module._parameters[tensor_name].device): # pylint: disable=protected-access + elif value is not None or not devices.same_device(device, module._parameters[tensor_name].device): # pylint: disable=protected-access param_cls = type(module._parameters[tensor_name]) # pylint: disable=protected-access module._parameters[tensor_name] = param_cls(value, requires_grad=old_value.requires_grad).to(device) # pylint: disable=protected-access t1 = time.time() diff --git a/modules/sd_offload.py b/modules/sd_offload.py index e36db4332..70777a30a 100644 --- a/modules/sd_offload.py +++ b/modules/sd_offload.py @@ -171,7 +171,7 @@ class OffloadHook(accelerate.hooks.ModelHook): return module def pre_forward(self, module, *args, **kwargs): - if devices.normalize_device(module.device) != devices.normalize_device(devices.device): + if not devices.same_device(module.device, devices.device): device_index = torch.device(devices.device).index if device_index is None: device_index = 0 From 42b3e08e658e5966e037c3faed38e6b117e64f3c Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 18:43:27 -0400 Subject: [PATCH 23/25] Control add setting to run hires with or without control Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 2 ++ modules/api/endpoints.py | 2 +- modules/control/units/controlnet.py | 2 +- modules/processing_diffusers.py | 6 +++--- modules/shared.py | 9 +++++---- modules/ui_control.py | 24 +++++++++++++++++++++++- 6 files changed, 35 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a1ed3e66a..90a7c43d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,8 @@ - Add [JoyCaption Beta](https://huggingface.co/fancyfeast/llama-joycaption-beta-one-hf-llava) support (in addition to existing JoyCaption Alpha) - Support Remote VAE with *Omnigen, Lumina 2 and PixArt* - Use Diffusers version of *OmniGen* + - Control move global settings to control elements -> control settings tab + - Control add setting to run hires with or without control - **SDNQ Quantization** - Add modules_to_not_convert support for post mode diff --git a/modules/api/endpoints.py b/modules/api/endpoints.py index 87181ba4f..561afafe0 100644 --- a/modules/api/endpoints.py +++ b/modules/api/endpoints.py @@ -25,7 +25,7 @@ def get_sd_models(): def get_controlnets(model_type: Optional[str] = None): from modules.control.units.controlnet import api_list_models - return api_list_models() + return api_list_models(model_type) def get_hypernetworks(): return [{"name": name, "path": shared.hypernetworks[name]} for name in shared.hypernetworks] diff --git a/modules/control/units/controlnet.py b/modules/control/units/controlnet.py index debe54c16..9233b47f8 100644 --- a/modules/control/units/controlnet.py +++ b/modules/control/units/controlnet.py @@ -152,7 +152,7 @@ def api_list_models(model_type: str = None): model_list += list(predefined_sd3) model_list += sorted(find_models()) return model_list - + def list_models(refresh=False): import modules.shared diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index 03d0b7b78..c3eb66d77 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -170,7 +170,7 @@ def process_hires(p: processing.StableDiffusionProcessing, output): prev_job = shared.state.job # hires runs on original pipeline - if hasattr(shared.sd_model, 'restore_pipeline') and shared.sd_model.restore_pipeline is not None: + if hasattr(shared.sd_model, 'restore_pipeline') and (shared.sd_model.restore_pipeline is not None) and not shared.opts.control_hires: shared.sd_model.restore_pipeline() # upscale @@ -200,8 +200,8 @@ def process_hires(p: processing.StableDiffusionProcessing, output): if 'Upscale' in shared.sd_model.__class__.__name__ or 'Flux' in shared.sd_model.__class__.__name__ or 'Kandinsky' in shared.sd_model.__class__.__name__: output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.width, height=p.height) if p.is_control and hasattr(p, 'task_args') and p.task_args.get('image', None) is not None: - if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0: - output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input + if hasattr(shared.sd_model, "vae") and output.images is not None and len(output.images) > 0: + output.images = processing_vae.vae_decode(latents=output.images, model=shared.sd_model, vae_type=p.vae_type, output_type='pil', width=p.hr_upscale_to_x, height=p.hr_upscale_to_y) # controlnet cannnot deal with latent input update_sampler(p, shared.sd_model, second_pass=True) orig_denoise = p.denoising_strength p.denoising_strength = strength diff --git a/modules/shared.py b/modules/shared.py index 276143e42..0f01696e9 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -895,10 +895,11 @@ options_templates.update(options_section(('postprocessing', "Postprocessing"), { })) options_templates.update(options_section(('control', "Control Options"), { - "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1}), - "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options"), - "control_move_processor": OptionInfo(False, "Processor move to CPU after use"), - "control_unload_processor": OptionInfo(False, "Processor unload after use"), + "control_hires": OptionInfo(False, "Use control during hires", gr.Checkbox, {"visible": False}), + "control_max_units": OptionInfo(4, "Maximum number of units", gr.Slider, {"minimum": 1, "maximum": 10, "step": 1, "visible": False}), + "control_tiles": OptionInfo("1x1, 1x2, 1x3, 1x4, 2x1, 2x1, 2x2, 2x3, 2x4, 3x1, 3x2, 3x3, 3x4, 4x1, 4x2, 4x3, 4x4", "Tiling options", gr.Textbox, {"visible": False}), + "control_move_processor": OptionInfo(False, "Processor move to CPU after use", gr.Checkbox, {"visible": False}), + "control_unload_processor": OptionInfo(False, "Processor unload after use", gr.Checkbox, {"visible": False}), })) options_templates.update(options_section(('interrogate', "Interrogate"), { diff --git a/modules/ui_control.py b/modules/ui_control.py index 94ad27eee..4b53a76bf 100644 --- a/modules/ui_control.py +++ b/modules/ui_control.py @@ -467,9 +467,31 @@ def create_ui(_blocks: gr.Blocks=None): if i == 0: units[-1].enabled = True # enable first unit in group - with gr.Accordion('Processor settings', open=False, elem_classes=['control-settings']) as _tab_settings: + with gr.Accordion('Control settings', open=False, elem_classes=['control-settings']) as _tab_settings: with gr.Group(elem_classes=['processor-group']): settings = [] + with gr.Accordion('Global', open=True, elem_classes=['processor-settings']): + control_hires = gr.Checkbox(label="Use control during hires", value=shared.opts.control_hires, elem_id='control_hires') + def set_control_hires(value): + shared.opts.control_active = value + control_hires.change(fn=set_control_hires, inputs=[control_hires], outputs=[]) + control_max_units = gr.Slider(label="Maximum units", minimum=1, maximum=10, step=1, value=shared.opts.control_max_units, elem_id='control_max_units') + def set_control_max_units(value): + shared.opts.control_max_units = value + control_max_units.change(fn=set_control_max_units, inputs=[control_max_units], outputs=[]) + control_tiles = gr.Textbox(label="Tiling options", value=shared.opts.control_tiles, elem_id='control_tiles') + def set_control_tiles(value): + shared.opts.control_tiles = value + control_tiles.change(fn=set_control_tiles, inputs=[control_tiles], outputs=[]) + control_move_processor = gr.Checkbox(label="Move processor to CPU after use", value=shared.opts.control_move_processor, elem_id='control_move_processor') + def set_control_move_processor(value): + shared.opts.control_move_processor = value + control_move_processor.change(fn=set_control_move_processor, inputs=[control_move_processor], outputs=[]) + control_unload_processor = gr.Checkbox(label="Unload processor after use", value=shared.opts.control_unload_processor, elem_id='control_unload_processor') + def set_control_unload_processor(value): + shared.opts.control_unload_processor = value + control_unload_processor.change(fn=set_control_unload_processor, inputs=[control_unload_processor], outputs=[]) + with gr.Accordion('HED', open=True, elem_classes=['processor-settings']): settings.append(gr.Checkbox(label="Scribble", value=False)) with gr.Accordion('Midas depth', open=True, elem_classes=['processor-settings']): From b6a0332694e3bf8cc43ba1bd58da016c08b17cb5 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Wed, 25 Jun 2025 18:45:40 -0400 Subject: [PATCH 24/25] update changelog Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 90a7c43d1..cd2156a3e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -35,6 +35,7 @@ - Fix delete file with gallery views - Add `SD_SAVE_DEBUG` env variable to report all params and metadata save operations as they happen - Fix TAESD model type detection + - Fix LoRA loader incorrectly reporting errors ## Update for 2025-06-16 From c25d15398fca39d5e085f12f3359a7c0b5c1fef5 Mon Sep 17 00:00:00 2001 From: Disty0 Date: Thu, 26 Jun 2025 01:48:01 +0300 Subject: [PATCH 25/25] SDNQ pre-mode don't quant if weights_dtype == 'none' --- modules/model_quant.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/modules/model_quant.py b/modules/model_quant.py index c8fd0f9b7..307707708 100644 --- a/modules/model_quant.py +++ b/modules/model_quant.py @@ -125,6 +125,8 @@ def create_sdnq_config(kwargs = None, allow_sdnq: bool = True, module: str = 'Mo return kwargs else: weights_dtype = shared.opts.sdnq_quantize_weights_mode + if weights_dtype is None or weights_dtype == 'none': + return kwargs if shared.opts.device_map == "gpu": quantization_device = devices.device