From 4db62cc7715dee0c86a6e565368b7a3675718280 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 16 May 2026 08:27:11 +0200 Subject: [PATCH] update xet and patch qk_norm defaults Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 3 +++ modules/models_hf.py | 23 ++++++++++------------- modules/sd_hijack_hfhub.py | 8 ++++---- modules/ui_definitions.py | 2 +- package.json | 2 +- pipelines/bria/transformer_block.py | 4 +++- pipelines/meissonic/transformer.py | 4 +++- pipelines/vibe/vibe_sana_editing.py | 18 ++++++++++++++---- requirements.txt | 3 ++- 9 files changed, 41 insertions(+), 26 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1cd2860ee..d49041554 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,8 @@ - **Captioning** new feature: analyze existing images for prompt adherence *tip*: image analysis requires larger VLM model to produce quality output new api endpoint: `/sdapi/v1/analyze` + - **HF download** use `XET` by default + see *settings -> huggingface -> download method* for options - **AI** - Cognitive analysis and improvements to *all* AI prompts - Automated fixes using `/check-` skills @@ -18,6 +20,7 @@ - `gradio` initial hijack - `SmolVLM` captioning - `gradio` temp files guard against large image + - `diffusers` patch custom pipelines for `qk_norm` ## Update for 2026-05-13 diff --git a/modules/models_hf.py b/modules/models_hf.py index c3e643b06..a5db96191 100644 --- a/modules/models_hf.py +++ b/modules/models_hf.py @@ -1,7 +1,6 @@ import os import time import gradio as gr -from installer import install from modules.logger import log from modules.shared import opts @@ -19,22 +18,20 @@ def hf_init(): os.environ.setdefault('HF_ENABLE_PARALLEL_LOADING', 'true' if opts.sd_parallel_load else 'false') os.environ.setdefault('HF_HUB_CACHE', opts.hfcache_dir) os.environ.setdefault('HF_XET_CACHE', opts.xetcache_dir) - if opts.hf_transfer_mode == 'requests': + if opts.hf_transfer_mode == 'HTTP': os.environ.setdefault('HF_XET_HIGH_PERFORMANCE', 'false') - os.environ.setdefault('HF_HUB_ENABLE_HF_TRANSFER', 'false') os.environ.setdefault('HF_HUB_DISABLE_XET', 'true') - elif opts.hf_transfer_mode == 'rust': - install('hf_transfer') + elif opts.hf_transfer_mode == 'XET': os.environ.setdefault('HF_XET_HIGH_PERFORMANCE', 'false') - os.environ.setdefault('HF_HUB_ENABLE_HF_TRANSFER', 'true') - os.environ.setdefault('HF_HUB_DISABLE_XET', 'true') - elif opts.hf_transfer_mode == 'xet': - install('hf_xet') - import huggingface_hub - huggingface_hub.utils._runtime.is_xet_available = lambda: True # pylint: disable=protected-access - os.environ.setdefault('HF_XET_HIGH_PERFORMANCE', 'true') - os.environ.setdefault('HF_HUB_ENABLE_HF_TRANSFER', 'true') os.environ.setdefault('HF_HUB_DISABLE_XET', 'false') + elif opts.hf_transfer_mode == 'XET HighPerformance': + os.environ.setdefault('HF_XET_HIGH_PERFORMANCE', 'true') + os.environ.setdefault('HF_HUB_DISABLE_XET', 'false') + os.environ.setdefault('HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY', 'false') + elif opts.hf_transfer_mode == 'XET Sequential': + os.environ.setdefault('HF_XET_HIGH_PERFORMANCE', 'false') + os.environ.setdefault('HF_HUB_DISABLE_XET', 'false') + os.environ.setdefault('HF_XET_RECONSTRUCT_WRITE_SEQUENTIALLY', 'true') obfuscated_token = None if len(opts.huggingface_token) > 0 and opts.huggingface_token.startswith('hf_'): diff --git a/modules/sd_hijack_hfhub.py b/modules/sd_hijack_hfhub.py index 17be5a172..e200bfb63 100644 --- a/modules/sd_hijack_hfhub.py +++ b/modules/sd_hijack_hfhub.py @@ -9,14 +9,14 @@ orig_xet_get = None def http_get_hijack(*args, **kwargs): - from modules.shared import state + from modules.shared import state, opts if len(args) > 0 and isinstance(args[0], str) and args[0].endswith(".json"): return orig_http_get(*args, **kwargs) jobid = state.begin('Download') fn = kwargs.get("displayed_filename", None) size = kwargs.get("expected_size", None) if fn and not fn.endswith(".json") and size is not None and size > 10240: - log.debug(f'Download: type=http fn="{fn}" size={size}') + log.debug(f'Download: type=http mode="{opts.hf_transfer_mode}" fn="{fn}" size={size}') debug(f'Download start: type=http args={args} kwargs={kwargs}') t0 = time.time() res = orig_http_get(*args, **kwargs) @@ -27,14 +27,14 @@ def http_get_hijack(*args, **kwargs): def xet_get_hijack(*args, **kwargs): - from modules.shared import state + from modules.shared import state, opts if len(args) > 0 and isinstance(args[0], str) and args[0].endswith(".json"): return orig_xet_get(*args, **kwargs) jobid = state.begin('Download') fn = kwargs.get("displayed_filename", None) size = kwargs.get("expected_size", None) if fn and not fn.endswith(".json"): - log.debug(f'Download: type=xet fn="{fn}" size={size}') + log.debug(f'Download: type=xet mode="{opts.hf_transfer_mode}" fn="{fn}" size={size}') debug(f'Download start: type=xet args={args} kwargs={kwargs}') res = orig_xet_get(*args, **kwargs) debug(f'Download end: type=xet res={res}') diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py index 889901119..86607bdf5 100644 --- a/modules/ui_definitions.py +++ b/modules/ui_definitions.py @@ -606,7 +606,7 @@ def create_settings(cmd_opts): "huggingface_sep": OptionInfo("

Huggingface

", "", gr.HTML), "diffuser_cache_config": OptionInfo(True, "Use cached model config when available"), "huggingface_token": OptionInfo('', 'HuggingFace token', gr.Textbox, {"lines": 2}, secret=True, env_var='HF_TOKEN'), - "hf_transfer_mode": OptionInfo("rust", "HuggingFace download method", gr.Radio, {"choices": ['requests', 'rust', 'xet']}), + "hf_transfer_mode": OptionInfo("XET", "HuggingFace download method", gr.Dropdown, {"choices": ['HTTP', 'XET', 'XET HighPerformance', 'XET Sequential']}), "huggingface_mirror": OptionInfo('', 'HuggingFace mirror', gr.Textbox), "offline_mode": OptionInfo(False, 'Force offline mode', gr.Checkbox), diff --git a/package.json b/package.json index b280e6231..38520a29b 100644 --- a/package.json +++ b/package.json @@ -18,7 +18,7 @@ "venv": ". venv/bin/activate", "start": ". venv/bin/activate; python launch.py --debug", "localize": "node cli/localize.js", - "packages": ". venv/bin/activate && pip install --upgrade accelerate huggingface_hub safetensors tokenizers peft pytorch_lightning pylint ruff", + "packages": ". venv/bin/activate && pip install --upgrade accelerate huggingface_hub hf_xet safetensors tokenizers peft pytorch_lightning pylint ruff", "format": ". venv/bin/activate && pre-commit run --all-files", "format-win": "venv\\scripts\\activate && pre-commit run --all-files", "eslint": "eslint javascript/", diff --git a/pipelines/bria/transformer_block.py b/pipelines/bria/transformer_block.py index 78dc3bc58..d909b6d83 100644 --- a/pipelines/bria/transformer_block.py +++ b/pipelines/bria/transformer_block.py @@ -106,9 +106,11 @@ class FluxSingleTransformerBlock(nn.Module): @maybe_allow_in_graph class FluxTransformerBlock(nn.Module): def __init__( - self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6 + self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str | bool | None = "rms_norm", eps: float = 1e-6 ): super().__init__() + if isinstance(qk_norm, bool): + qk_norm = "rms_norm" if qk_norm else None self.norm1 = AdaLayerNormZero(dim) self.norm1_context = AdaLayerNormZero(dim) diff --git a/pipelines/meissonic/transformer.py b/pipelines/meissonic/transformer.py index c2336d323..274cbb7cb 100644 --- a/pipelines/meissonic/transformer.py +++ b/pipelines/meissonic/transformer.py @@ -440,8 +440,10 @@ class TransformerBlock(nn.Module): processing of `context` conditions. """ - def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm="rms_norm", eps=1e-6): + def __init__(self, dim, num_attention_heads, attention_head_dim, qk_norm: str | bool | None = "rms_norm", eps=1e-6): super().__init__() + if isinstance(qk_norm, bool): + qk_norm = "rms_norm" if qk_norm else None self.norm1 = AdaLayerNormZero(dim) diff --git a/pipelines/vibe/vibe_sana_editing.py b/pipelines/vibe/vibe_sana_editing.py index 3b88bed2e..4f013a3dc 100644 --- a/pipelines/vibe/vibe_sana_editing.py +++ b/pipelines/vibe/vibe_sana_editing.py @@ -23,6 +23,14 @@ from .edit_head import MetaConnector logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def _normalize_qk_norm(qk_norm: str | bool | None) -> str | None: + """Normalize legacy boolean qk_norm values to Diffusers-compatible strings.""" + if isinstance(qk_norm, bool): + # Legacy checkpoints can store qk_norm as bool while current diffusers expects a string. + return "rms_norm" if qk_norm else None + return qk_norm + + class SanaLinearAttnProcessor2_0: r"""Processor for implementing scaled dot-product linear attention.""" @@ -112,7 +120,7 @@ class SanaTransformerBlock(nn.Module): cross_attention_dim: int = 2240, norm_eps: float = 1e-6, mlp_ratio: float = 2.5, - qk_norm: str | None = None, + qk_norm: str | bool | None = None, *, attention_out_bias: bool = True, attention_bias: bool = True, @@ -130,12 +138,13 @@ class SanaTransformerBlock(nn.Module): cross_attention_dim (int): The dimension of the cross-attention context. norm_eps (float): Epsilon value for layer normalization stability. mlp_ratio (float): Expansion ratio for the feed-forward network hidden dimension. - qk_norm (str | None): Normalization method for Query/Key vectors. + qk_norm (str | bool | None): Normalization method for Query/Key vectors. attention_out_bias (bool): Whether to include bias in the attention output projection. attention_bias (bool): Whether to include bias in the attention Q/K/V projections. norm_elementwise_affine (bool): Whether to learn affine parameters for normalization. """ super().__init__() + qk_norm = _normalize_qk_norm(qk_norm) # 1. Self Attention self.norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=norm_eps) @@ -269,7 +278,7 @@ class VIBESanaEditingModel(SanaTransformer2DModel): patch_size: int = 1, norm_eps: float = 1e-6, interpolation_scale: int | None = None, - qk_norm: str | None = None, + qk_norm: str | bool | None = None, timestep_scale: float = 1.0, input_condition_type: str = "channel_cat", edit_head_input_dim: int = 2048, @@ -299,7 +308,7 @@ class VIBESanaEditingModel(SanaTransformer2DModel): patch_size (int): Size of the patches extracted from the input latent. norm_eps (float): Epsilon for layer normalization. interpolation_scale (int | None): Scale factor for positional embedding interpolation. - qk_norm (str | None): Normalization type for Query/Key (e.g., 'rms_norm'). + qk_norm (str | bool | None): Normalization type for Query/Key (e.g., 'rms_norm'). timestep_scale (float): Scale factor for the timestep. input_condition_type (str): Method for conditioning on the input image. Options: 'channel_cat', 'seq_cat'. edit_head_input_dim (int): Input dimension for the `MetaConnector` edit head. @@ -311,6 +320,7 @@ class VIBESanaEditingModel(SanaTransformer2DModel): guidance_embeds (bool): Whether to use additional guidance embeddings. """ super(SanaTransformer2DModel, self).__init__() + qk_norm = _normalize_qk_norm(qk_norm) out_channels = out_channels or in_channels inner_dim = num_attention_heads * attention_head_dim diff --git a/requirements.txt b/requirements.txt index f1391e86a..c315c9e86 100644 --- a/requirements.txt +++ b/requirements.txt @@ -30,7 +30,8 @@ requests==2.32.3 tqdm==4.67.3 accelerate==1.13.0 einops==0.8.1 -huggingface_hub==1.14.0 +huggingface_hub==1.15.0 +hf_xet==1.5.0 numpy==2.1.2 pandas==2.3.1 protobuf==6.33.5