Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-18 08:44:33 +02:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #5086 from vladmandic/dev

Browse Source 
dev merge
This commit is contained in:
Vladimir Mandic
2026-09-17 07:53:47 +02:00
committed by GitHub
parent 684940e015 dca0bc6924
commit 112d5afbfb
297 changed files with 26112 additions and 6128 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+160 -3
View File
@@ -1,10 +1,167 @@
# Change Log for SD.Next
## Highlights for 2026-08-26
## Update for 2026-09-17
### Highlights for 2026-09-17
*What's New*? Well, code-wise, this is a big one...
First, a-lot-of-optimizations:
- Updated core packages
- Improved **LoRA** performance and quality, especially with quantized models
- Newly structured **attention** mechanisms
- Modular pipelines with new **guidance** methods
- Support for different **caching** stacks
- Compute updates across the board
And some cool new stuff and models:
- **DLSS v5** integration
- New models: **Anima 2.9B**, **LLaDa-Image**
- And few cloud models: *Google's Gemini, NanoBanana, Veo, Omni* and *X.AI's Grok*
- Some (light) UI restyling
Plus inevitable bug-fixes...
[Home](https://vladmandic.github.io/sdnext/) | [ChangeLog](https://github.com/vladmandic/automatic/blob/master/CHANGELOG.md) | [Docs](https://vladmandic.github.io/sdnext-docs/) | [Discord](https://discord.com/invite/sd-next-federal-batch-inspectors-1101998836328697867) | [Sponsor](https://github.com/sponsors/vladmandic)
### Details for 2026-09-17
- **Models**
- [Anima 2.9B Preview v1](https://huggingface.co/yeoj34760/Anima-2.9B)
expanded version of Anima 2B
- [inclusionAI LLaDA-Image](https://huggingface.co/inclusionAI/LLaDA-Image) in *base* and *turbo* variants
LLaDA-Image is a 6.5B transformer with massive 16.3B fully-custom MoE text-encoder and optional 1.3B SigVQ conditioning model
with support for text-to-image, vq-conditioned text-to-image and image-editing workflows
*note* model is extremely quantization sensitive so minimum allowed quant type is `uint8`
- [MiniMax-H3](https://huggingface.co/MiniMaxAI/MiniMax-H3) updates
new [SDNQ-uint8](https://huggingface.co/OzzyGT/MiniMax_H3_sdnq_8bit_pruned) pre-quantized *base* and *pruned* variants
new [Nunchaku-Lite](https://huggingface.co/rootonchair/MiniMax-H3-nunchaku-lite-int4) variant
new [VDN](https://huggingface.co/OpenVDN/vdn-minimax-h3) *video-delta-net* variant
- **LoRA**
- see [LoRA docs](https://vladmandic.github.io/sdnext-docs/LoRA) for all of the improvements and usage instructions
*note*: lora now has its own settings section in *settings -> lora*
- new apply engine that allows lora to be applied much faster
- new calibration engine that allows lora to be applied with far smaller error when dealing with highly quantized models
- *note*: calibration data is stored once calculated so it can be reused for future runs
location is `models/calibration` folder
- new factor cache that allows lora effects to be pre-calculated and persistently cached for future runs
location is `models/lora-factor-cache` folder
- multi-network stack modes
can significantly improve lora quality when using multiple loras at once
- per-block strength
- native support for **MiniMax**
see [MiniMax Turbo LoRA collection](https://huggingface.co/vladmandic/MiniMax-H3-Turbo-LoRA) for LoRAs and examples
- **DLSS**
- add DLSS support for: *NeuralRender, SuperSample and FrameGen*
dlls 5 caused quite a stir, but combined with generative ai it becomes a nice tool
- available as part of image/video generate workflows via *extras -> dlss*
or as a standalone *processing* workflow
or via xyz grid
- *note*: requires nvidia rtx gpu, windows platform and compatible gpu drivers
but...it can be used from wsl2: unpack required package on windows host and you can access it from the wsl2 environment
- *install*: requires [DLSS 5 Visual Enhancer](https://github.com/Merserk/dlss5-visual-enhancer/releases/tag/v7.0)
- *diag*: enable `SD_DLSS_DEBUG=true` and monitor `dlss.log` in the package directory
- **Attention**
- see [Attention docs](https://vladmandic.github.io/sdnext-docs/Attention) for details and usage instructions
*note*: attention now has its own settings section in *settings -> cross attention*
*note*: this is a breaking change - if you had custom attention settings in previous releases, you will need to re-apply them in the new settings section
- new `sparse-attention` method that can be combined with other attention methods
to reduce memory usage and improve performance on large models
- new attention mechanisms decision tree and apply method refactor
- **Modular Pipelines**
- see [Modular Pipelines docs](https://vladmandic.github.io/sdnext-docs/Modular-Pipelines) for details and usage instructions
- new model **Guidance** stack for modular pipelines
includes: *CFG, PAG, Auto, Zero, APG, SLG, SEG, TCFG, FDG*
see [Guidance docs](https://vladmandic.github.io/sdnext-docs/Guidance) for details and usage instructions
- new model **Caching** stack for modular pipelines
includes: *FasterCache, FirstBlockCache, LayerSkip, MagCache, PyramidAttentionBroadcast, TaylorSeerCache, TextKVCache*
see [Caching docs](https://vladmandic.github.io/sdnext-docs/Caching) for details and usage instructions
- implement progress and preview
- intercept and profiling hooks
- on-demand convert standard model on-demand
- **Cloud**
- updated support for google models in text, image and video workflows
*note*: requires google api key
- [Google Veo](https://ai.google.dev/gemini-api/docs/veo) in *preview*, *fast* and *lite* variants
workflows: *t2v, i2v*
- [Google Omni](https://ai.google.dev/gemini-api/docs/omni) in *flash* variant
workflows: *t2v, i2v*
- [Google Nano Banana](https://ai.google.dev/gemini-api/docs/models/gemini-3.1-flash-image) in *2* and *2 lite* and *pro* variants
workflows: *caption*
- [Google Gemini](https://ai.google.dev/gemini-api/docs/models/gemini-3.8-flash) in *flash* and *pro* variants
workflows: *caption, prompt-enhance*
- added support for xai grok models
- [X.AI Grok](https://x.ai/grok) in *3*, *3 fast*, *3 mini* and *3 mini fast* variants
workflows: *caption, prompt-enhance*
*note*: requires grok api key
- **Compute**
- cuda: update `torch==2.14.0` with `cuda==13.2`
- openvino: update `openvino==2026.3.1` with `torch==2.13.0`
- option to skip triton autotune and use default config for all triton kernels
in *settings -> compute settings*
*note*: this may improve initial generate time, but may also reduce performance on some models
- rocm: update `rocm` script and add detailed `miopen` logging, thanks @resonantsky
- new optional transformer hooks
in *settings -> compute add-ons*
*PAG: Perturbed attention guidance, PAB: Pyramid attention broadcast, FBC: First Block Cache, FC: Faster Cache, LS: Layer Skip, MC: Mag Cache, TS: TaylorSeer*
*note*: compatibility of different methods varies across different models
- update `numpy` and `scipy` frozen requirements as required by new compute drivers
*note*: this may break compatibility with some legacy packages, so report any finidings
- **Other**
- video preview: TAESD support for **MiniMax**
- support `xai grok` for prompt enhance workflows
*note*: requires grok api key
- remove `/redocs` as `/docs` are primary api docs
- rebuild docs site index
- **UI**
- some (light) re-styling of the *Default* theme
- add new *Tillerz-CleanDark* theme, thanks @Tillerz
- ability to filter samplers and upscalers, thanks @emecii
- **Wiki/Docs**:
- new articles: *Attention, Modular-Pipelines*
- updated: *LoRA, MiniMax*
- **Fixes**
- api: prompt enhance with vision
- autocomplete: skip disabled networks
- compile: keep model compiled state
- detailer: handling of stop/skip/pause
- framepack: correct device assignment, thanks @li-lizhe
- group offload: improve memory management
- installer: better handle git detached head
- json: handle file locks
- log: ansi color handling
- lora: cleanup tags
- lora: support transformer ref models
- lora: keep parsed network data through pipeline
- lucida: handle requirements
- lumina-dimoo: attention-kwargs, thanks @Anai-Guo
- metadata: fix wildcard info
- minimax: crop image to video aspect ratio
- modular: handle module with remote-code
- network: improve type/version lookup
- offline: honor offline mode for more models, thanks @ryanmeador
- openvino: optimize recompile checks and lora loading
- prompt enhance: cloud models use correct system prompt
- prompt enhance: use init image for video
- prompt: cache checks when cfg changes
- prompt: unnecessary secondary prompt if same
- prompt: clean prompt after network parsing
- rife: cleanup dead code, thanks @Anai-Guo
- temp files: handle locking
- theme: fix circular imports changing theme to default
- todo: remove dead code, thanks @Anai-Guo
- ui: js fetch exception handling
- update: handle git errors gracefully
- vae: fetch scale factor from the model
- vdm scheduler: fix steps, thanks @zjn20030811
- xyz grid: apply bool values
## Update for 2026-08-26
### Highlights for 2026-08-26
Time for a new release, *this is a large one*!
Main focus is improving video workflows which also brings full support for new [MiniMax H3](https://vladmandic.github.io/sdnext-docs/MiniMax) and [LTXVideo-2.5](https://vladmandic.github.io/sdnext-docs/LTX)
and improves general video processing with flexible video upscaling, updated interpolation, etc.
and improvements to general video processing with flexible video upscaling, updated interpolation, etc.
*What else?*
- [Detailer.next](https://vladmandic.github.io/sdnext-docs/Detailer) with new support for *vision-language models* and *per-class prompts*
@@ -1898,7 +2055,7 @@ And check out new **history** tab in the right panel, it now shows visualization
*note*: this does not impact the actual image resolution, only the resolution at which detailer internally operates
- refactor reuse-seed and add functionality to all tabs
- refactor modernui js codebase
- move zluda flash attenion to *Triton Flash attention* option
- move zluda flash attenion to *Triton AMD Flash attention* option
- remove samplers filtering
- allow both flow-matching and discrete samplers for sdxl models
- cleanup command line parameters
TODO.md
+11 -24
View File
@@ -2,15 +2,15 @@
## Short-term
- LoRA: merge new handler, @CalamitousFelicitousness
- Attn: merge refactor, @CalamitousFelicitousness
- MiniMax LoRA: native loader for MiniMax-H3: fl2va, ref2va, pruned
- MiniMax TAESD: <https://github.com/madebyollin/taehv>
- MiniMax: Create pre-quant for MiniMax-H3-Turbo
- Benchmark tool productize: @CalamitousFelicitousness
- Inpaint: https://discord.com/channels/1101998836328697867/1130536562422186044/1506850651035144322, @vladmandic
- Control tab verify overrides handling, @vladmandic
- LTX: Create pre-quant for LTX-2.5
- LTX: Implement LTX2DFRPipeline
- ROCm: v10
- Video: unify execution path for ui and api
- Torch: update ipex, rocm to torch==2.14
## Issues
- [Inpaint](https://discord.com/channels/1101998836328697867/1130536562422186044/1506850651035144322), @vladmandic
## Features
@@ -21,20 +21,19 @@
- Lightweight scheduler/queue manager, @vladmandic
- Integrate natural language image search: [ImageDB](https://github.com/vladmandic/imagedb), @vladmandic
- Support cloud providers, @CalamitousFelicitousness
- Benchmark tool productize: @CalamitousFelicitousness
### Roadmap
- Automated testing and integrate models repo
- Video upscaling: LTX-Upscaler
- Video upscaling: [MiniMax-Upscaler](https://huggingface.co/LBH-123-AI/Minimax_h3_latent_Upscaler)
- Video capabilities to processing tab, add RIFE, upscaling (once available)
- Distraction-free UI mode with prompt-only, chat-based interface
- Revisit transformer caching for modular pipelines
- Revisit guidance for modular pipelines
- Implement modular for some image models
- Video models: support finetunes
- Incorporate [prompting guides](https://github.com/CalamitousFelicitousness/ai-prompting-guides)
- Video models: use Networks/Reference instead of custom
- UI Lite vs Expert mode
- Expand custom VAE support
- Remove obsolete code: `olive-ai`
### OnHold
@@ -51,18 +50,6 @@
- Unify *huggingface* and *diffusers* model folders
- JSON image metadata
### Modular
*Pending finalization of modular pipelines implementation and development of compatibility layer*
- Switch to modular pipelines
- Feature: Transformers unified cache handler
- Refactor: [Modular pipelines and guiders](https://github.com/huggingface/diffusers/issues/11915)
- [MagCache](https://github.com/huggingface/diffusers/pull/12744)
- [SmoothCache](https://github.com/huggingface/diffusers/issues/11135)
- [STG](https://github.com/huggingface/diffusers/blob/main/examples/community/README.md#spatiotemporal-skip-guidance)
- [TextKVCache](https://huggingface.co/NucleusAI/Nucleus-Image#quick-start), @vladmandic
## New models / Pipelines
TODO: Investigate which models are diffusers-compatible and prioritize!
cli/api-enhance.py
+2 -1
View File
@@ -59,6 +59,7 @@ def enhance(args): # pylint: disable=redefined-outer-name
options['model'] = str(args.model)
if args.image:
options['image'] = encode(args.image)
options['use_vision'] = True
response = post('/sdapi/v1/prompt-enhance', options)
return response
@@ -72,6 +73,6 @@ if __name__ == "__main__":
parser.add_argument('--image', type=str, default=None, required=False, help='optional input image')
parser.add_argument('--nsfw', type=bool, action=argparse.BooleanOptionalAction, required=False, help='nsfw allowed')
args = parser.parse_args()
log.info(f'api-upscale: {args}')
log.info(f'api-enhance: {args}')
result = enhance(args)
log.info(result)
cli/lora-quant-fidelity.py
+685
View File
@@ -0,0 +1,685 @@
#!/usr/bin/env python
"""LoRA fidelity analyzer for quantized base models.
Measures, in weight space, how faithfully a LoRA lands on an SDNQ-quantized
model. Every targeted module is rebuilt with the loader's own module class and
its delta taken from the production ``calc_updown``, so all adapter families
(LoRA, LoKR, LoHA, OFT, full, IA3, GLoRA, norm, plus DoRA and bias variants)
are measured as they would actually apply:
- factor path (plain additive LoRA riding the svd side-channel): storage is
lossless; the reported figure is the delta realized through the result-dtype
materialize, the same bf16 rounding an unquantized model applies.
Eligibility is decided by the loader's own predicate.
- hosted path (non-factorable families on sub-8-bit formats): the seeded svd
truncation at ``--host-rank``, realized the same way.
- requantize path (all other fallbacks): retention ``rho`` of the intended
delta. On-grid rounding erases sub-step deltas down to a ``2/group_size``
floor, so low-bit formats (<=6 bits) typically show rho ~= 0.02-0.03.
- unquantized modules: the LoRA applies exactly regardless.
Reported fidelity is per-module ``applied_rho`` (the measured figure for
whichever path the loader would take), summarized as a median and an
energy-weighted mean over the file's modules; ``requant_rho`` always carries
the if-merged figure. ``snr`` sets the delta against the base weight's own
quantization error (uniform rounding from the grid step, or measured with
``--reference``), in the calibrated norm with ``--calib``: it does not depend
on the apply path and compares checkpoints of different widths.
Works offline against a pre-quantized SDNQ repo (stored tensors + config,
streamed one module at a time so the repo need not fit in memory) or a bf16
repo with simulated quantization settings, so a combination can be assessed
before committing to a quantized checkpoint.
Examples:
python cli/lora-quant-fidelity.py --model vladmandic/Krea-2-Base-sdnq-hadamard-uint4 --arch krea2 --lora "~/models/Lora/Krea 2/krea2_turbo_distill_r256.safetensors"
python cli/lora-quant-fidelity.py --model CalamitousFelicitousness/Krea-2-Base-Diffusers --arch krea2 --dtype uint4 --lora lora.safetensors --json report.json
"""
import os
import re
import sys
import json
import types
import argparse
import importlib
import importlib.util
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
os.environ.setdefault('SD_INSTALL_QUIET', '1')
def parse_cli():
parser = argparse.ArgumentParser(description='lora-quant-fidelity')
parser.add_argument('--model', required=True, help='model dir, transformer dir, or org/name repo id')
parser.add_argument('--arch', default='generic', help='lora key resolver: a native arch (e.g. krea2, zimage, f2) or generic')
parser.add_argument('--lora', required=True, nargs='+', help='lora safetensors file(s)')
parser.add_argument('--dtype', default=None, help='simulate quantization of a bf16 repo at this sdnq dtype (e.g. uint4, int8); bf16 measures the unquantized reference')
parser.add_argument('--group', type=int, default=0, help='sdnq group_size for simulation')
parser.add_argument('--hadamard-group', type=int, default=256, help='sdnq hadamard group for simulation')
parser.add_argument('--sample', type=int, default=40, help='max modules analyzed per lora (evenly sampled)')
parser.add_argument('--full', action='store_true', help='analyze every matched module')
parser.add_argument('--json', default=None, help='write full report to this json file')
parser.add_argument('--host-rank', type=int, default=256, help='svd hosting cap for non-factorable modules on sub-8-bit formats, mirroring lora_sdnq_host_rank; 0 scores the requantize path instead')
parser.add_argument('--calib', default=None, help='activation statistics file (models/calibration/*.safetensors): hosting truncation is then channel-weighted as with lora_sdnq_host_calib, and hosted rho is measured in the activation-weighted norm (the output-error proxy)')
parser.add_argument('--reference', default=None, help='unquantized repo of the same model: the base quantization error each delta competes with is then measured per module instead of estimated from the grid step')
parser.add_argument('--fail-under', type=float, default=None, help='exit 2 when median applied fidelity of any lora is below this')
return parser.parse_args()
cli_args = parse_cli()
sys.argv = [sys.argv[0]] # sdnext arg parsing during imports must not see tool args (prefix matching eats --model/--lora)
import modules.cmd_args # pylint: disable=wrong-import-position
import installer # pylint: disable=wrong-import-position
modules.cmd_args.parse_args()
installer.add_args(modules.cmd_args.parser)
modules.cmd_args.parsed, _unknown = modules.cmd_args.parser.parse_known_args([])
import torch # pylint: disable=wrong-import-position
from safetensors import safe_open # pylint: disable=wrong-import-position
from rich import print as rprint # pylint: disable=wrong-import-position
from modules import shared # pylint: disable=wrong-import-position,unused-import # shared must initialize before sd_models, which imports back into it
from modules.lora import native_adapter, network, network_lora, network_lokr, network_hada, network_oft, network_full, network_ia3, network_glora, network_norm, lora_sdnq # pylint: disable=wrong-import-position
from modules.lora.lora_load import NATIVE_DISPATCH # pylint: disable=wrong-import-position
from sdnq.quantizer import sdnq_quantize_layer_weight # pylint: disable=wrong-import-position
from sdnq.quant_utils import rotate_hadamard # pylint: disable=wrong-import-position
MODEL_ROOTS = [
os.path.expanduser('~/database/models/huggingface'),
os.path.expanduser('~/database/models/Diffusers'),
]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ARCH_PACKAGES = {'zimage': 'z_image', 'f2': 'flux', 'minimaxh3': 'minimax'} # arches whose pipelines package is not spelled like the arch
# every adapter family the native loader can build, with the module class that owns its
# apply-time math. deltas are taken from the production calc_updown so the tool cannot
# drift from the loader, and eligibility is decided by the production predicate itself.
FAMILY_SPECS = (
('lora', network_lora.NetworkModuleLora, native_adapter.LORA_SUFFIXES, native_adapter.LORA_MARKERS),
('lokr', network_lokr.NetworkModuleLokr, native_adapter.LOKR_SUFFIXES, native_adapter.LOKR_MARKERS),
('loha', network_hada.NetworkModuleHada, native_adapter.LOHA_SUFFIXES, native_adapter.LOHA_MARKERS),
('oft', network_oft.NetworkModuleOFT, native_adapter.OFT_SUFFIXES, native_adapter.OFT_MARKERS),
('full', network_full.NetworkModuleFull, native_adapter.FULL_SUFFIXES, native_adapter.FULL_MARKERS),
('ia3', network_ia3.NetworkModuleIa3, native_adapter.IA3_SUFFIXES, native_adapter.IA3_MARKERS),
('glora', network_glora.NetworkModuleGLora, native_adapter.GLORA_SUFFIXES, native_adapter.GLORA_MARKERS),
('norm', network_norm.NetworkModuleNorm, native_adapter.NORM_SUFFIXES, native_adapter.NORM_MARKERS),
)
class StubOnDisk:
def __init__(self, path):
self.filename = path
self.name = os.path.splitext(os.path.basename(path))[0]
self.shorthash = ''
self.sd_version = 'unknown'
with safe_open(path, framework='pt', device='cpu') as f:
self.metadata = f.metadata() or {}
def resolve_model_dir(spec):
"""Return the transformer directory for a local path or org/name repo id."""
candidates = [spec, os.path.join(spec, 'transformer')]
cache_name = 'models--' + spec.replace('/', '--')
for root in MODEL_ROOTS:
snap_root = os.path.join(root, cache_name, 'snapshots')
if os.path.isdir(snap_root):
for snap in sorted(os.listdir(snap_root), reverse=True):
candidates.append(os.path.join(snap_root, snap, 'transformer'))
candidates.append(os.path.join(snap_root, snap))
for c in candidates:
if os.path.isfile(os.path.join(c, 'config.json')):
return c
raise SystemExit(f'model not found: {spec}')
def resolve_arch(name):
"""Return the arch lora module for key resolution, or None for generic matching."""
if name == 'generic':
return None
path = NATIVE_DISPATCH.get({'flux2': 'f2', 'ernie': 'ernieimage'}.get(name, name))
if path is None:
raise SystemExit(f'unknown arch {name}; choices: {sorted(NATIVE_DISPATCH)} or generic')
return importlib.import_module(path)
def map_lora_modules(lora_path, arch_mod):
"""Return {model_module_path: (family, weights)} across every adapter family, plus a census.
Grouping mirrors the native loader: a family is only considered when its
marker is present, and groups resolve to model paths through the arch's own
resolver. A fused save is sliced onto its targets for the lora family; the
other families' chunks are counted but not analyzed (their apply-time math
is arch-owned).
"""
with safe_open(lora_path, framework='pt', device='cpu') as f:
state_dict = {k: f.get_tensor(k) for k in f.keys()}
metadata = f.metadata() or {}
prefixes = getattr(arch_mod, 'KNOWN_PREFIXES', native_adapter.KNOWN_PREFIXES_DEFAULT)
resolve = getattr(arch_mod, 'resolve_targets', None) or (lambda prefix, base: [(base, None)])
grouper = getattr(arch_mod, 'group_by_suffixes', None) or native_adapter.group_by_suffixes # an arch that rewrites keys before parsing groups them itself
file_alpha = getattr(arch_mod, 'file_alpha', None)
mapped, census, chunked = {}, {}, 0
for fam, _cls, suffixes, markers in FAMILY_SPECS:
if not native_adapter.has_marker(state_dict, markers):
continue
groups = grouper(state_dict, suffixes, prefixes=prefixes)
if fam == 'lora':
groups = {k: w for k, w in groups.items() if 'lora_down.weight' in w and 'lora_up.weight' in w}
alpha = file_alpha(types.SimpleNamespace(filename=lora_path, metadata=metadata)) if file_alpha is not None else None
if alpha is not None and not any('alpha' in w for w in groups.values()): # a file-level alpha applies only to files without alpha tensors, as in try_load_lora
groups = {k: {**w, 'alpha': torch.tensor(float(alpha))} for k, w in groups.items()}
else:
groups = {k: w for k, w in groups.items() if native_adapter.has_marker({f'x.{s}': None for s in w}, markers)}
if not groups:
continue
census[fam] = len(groups)
for (prefix, base), w in groups.items():
for path, chunk in native_adapter.resolve_group_targets(resolve, prefix, base):
target = w
if chunk is not None: # a fused save spans several modules; only the lora family slices its up factor
target = None
if fam == 'lora':
fused_out = w['lora_up.weight'].shape[0]
target = native_adapter.slice_lora_chunk(w, chunk)
target = native_adapter.slice_dora_scale(target, chunk, fused_out)
target = native_adapter.slice_bias_delta(target, chunk, fused_out) if target is not None else None
if target is None:
chunked += 1
continue
mapped.setdefault(path, []).append((fam, target)) # a module can carry several families; the loader applies each
return mapped, census, chunked
def stamp_index(paths):
"""Map each module path to its stamped form, the way the loader matches.
The loader compares ``network_prefix + path.replace('.', '_')`` against each
module's stamped ``network_layer_name``, so kohya-style ``lora_unet_`` keys
(whose base arrives already underscored) resolve fine there. Matching on the
stamped form reproduces that and keeps dotted bases working unchanged.
"""
return {p.replace('.', '_'): p for p in paths}
def make_stub(shape, dtype=torch.bfloat16):
"""Minimal sd_module standing in for a bf16 repo weight: the module classes key off its type and shape."""
if len(shape) == 2:
return torch.nn.Linear(shape[1], shape[0], bias=False, dtype=dtype, device='meta')
return torch.nn.Conv2d(shape[1], shape[0], shape[2:], bias=False, dtype=dtype, device='meta')
def build_module(fam, path, w, net, sd_module):
"""Instantiate the family's production NetworkModule for one target."""
cls = next(c for f, c, _s, _m in FAMILY_SPECS if f == fam)
weights = network.NetworkWeights(network_key=path, sd_key=path, w=w, sd_module=sd_module)
return cls(net, weights)
def resolve_transformer_cls(arch, class_name, model_dir=None):
"""Resolve the transformer class: an sdnext-owned spec class, then diffusers, then a modeling file beside the weights.
Arches like krea2 keep checkpoint-style names in their own class; pruned MiniMax repos ship theirs as remote code."""
if not class_name:
return None
if arch:
try:
pkg = importlib.import_module(f'pipelines.{ARCH_PACKAGES.get(arch, arch)}')
for attr in dir(pkg):
if attr.endswith('_SPEC'):
cls = getattr(getattr(pkg, attr), 'cls', None)
if cls is not None and cls.__name__ == class_name:
return cls
except Exception:
pass
import diffusers
cls = getattr(diffusers, class_name, None)
if cls is not None or not model_dir:
return cls
for fname in sorted(os.listdir(model_dir)):
path = os.path.join(model_dir, fname)
if not fname.endswith('.py'):
continue
with open(path, encoding='utf-8') as f:
if re.search(rf'^class {re.escape(class_name)}\b', f.read(), re.MULTILINE) is None:
continue
spec = importlib.util.spec_from_file_location(fname[:-3], path)
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return getattr(module, class_name)
return None
class QuantRepo:
"""Per-module access to a pre-quantized SDNQ repo: a meta skeleton built through the sdnq conversion, each layer's
tensors streamed from the shards while it is analyzed; tensors outside the block stacks stay resident for the arch hooks."""
LAYER_KEYS = ('weight', 'bias', 'scale', 'zero_point', 'svd_up', 'svd_down')
def __init__(self, model_dir, model_config, arch=None):
from accelerate import init_empty_weights
from sdnq import SDNQConfig
from sdnq.quantizer import sdnq_post_load_quant
from sdnq.utils import get_quant_args_from_config
cls = resolve_transformer_cls(arch, model_config.get('_class_name'), model_dir)
if cls is None:
raise SystemExit(f'cannot resolve transformer class {model_config.get("_class_name")} for "{model_dir}"')
quant_config = SDNQConfig.from_dict(model_config['quantization_config'])
with init_empty_weights():
config = cls.load_config(model_dir)
if hasattr(config, 'pop'):
config.pop('quantization_config', None)
model = cls.from_config(config)
model = sdnq_post_load_quant(model, torch_dtype=torch.bfloat16, pre_quantized=True, **get_quant_args_from_config(quant_config))
self.model = model
self.model_dir = model_dir
self.handles = {}
index = os.path.join(model_dir, 'diffusion_pytorch_model.safetensors.index.json')
if os.path.isfile(index):
with open(index, encoding='utf-8') as f:
weight_map = json.load(f)['weight_map']
else:
weight_map = {}
for fname in sorted(os.listdir(model_dir)):
if fname.endswith('.safetensors'):
with safe_open(os.path.join(model_dir, fname), framework='pt', device='cpu') as f:
weight_map.update(dict.fromkeys(f.keys(), fname))
mapping = getattr(model, '_checkpoint_conversion_mapping', None) or {}
self.shards = {} # model key -> (stored key, shard)
for stored, shard in weight_map.items():
key = stored
for pattern, replacement in mapping.items():
key = re.sub(pattern, replacement, key)
self.shards[key] = (stored, shard)
self.layers = {}
for name, module in model.named_modules():
if getattr(module, 'sdnq_dequantizer', None) is not None or (module.__class__.__name__ == 'Linear' and getattr(module, 'weight', None) is not None):
self.layers[name] = module
stacks = tuple(f'{name}.' for name, module in model.named_modules() if isinstance(module, torch.nn.ModuleList))
resident = {key: self.get(key) for key in self.shards if not key.startswith(stacks)}
model.load_state_dict(resident, strict=False, assign=True)
def get(self, key):
entry = self.shards.get(key)
if entry is None:
return None
stored, shard = entry
f = self.handles.get(shard)
if f is None:
f = safe_open(os.path.join(self.model_dir, shard), framework='pt', device='cpu')
self.handles[shard] = f
return f.get_tensor(stored)
def materialize(self, name):
"""Load one layer's stored tensors onto the analysis device and return the layer."""
from sdnq.quant_utils import prepare_weight_for_matmul, prepare_svd_for_matmul
layer = self.layers[name]
state = {}
for local in self.LAYER_KEYS:
t = self.get(f'{name}.{local}')
if t is not None:
state[local] = t.to(device)
layer.load_state_dict(state, strict=False, assign=True)
deq = getattr(layer, 'sdnq_dequantizer', None)
if deq is not None: # the loader's post-processing, so the dequantizer sees the layout it expects
if deq.use_quantized_matmul and not deq.re_quantize_for_matmul:
layer.weight.data = prepare_weight_for_matmul(layer.weight, matmul_dtype=deq.quantized_matmul_dtype)
if getattr(layer, 'svd_up', None) is not None:
layer.svd_up.data, layer.svd_down.data = prepare_svd_for_matmul(layer.svd_up, layer.svd_down, deq.use_quantized_matmul)
return layer
def release(self, name):
layer = self.layers[name]
state = {local: torch.empty_like(t, device='meta') for local, t in layer.state_dict().items() if t.device.type != 'meta'}
layer.load_state_dict(state, strict=False, assign=True)
class Bf16Repo:
"""Lazy per-module weight access for a sharded bf16 transformer repo."""
def __init__(self, model_dir):
self.model_dir = model_dir
self.handles = {} # reopening a multi-gb shard per module dominates runtime over many loras
index = os.path.join(model_dir, 'diffusion_pytorch_model.safetensors.index.json')
if os.path.isfile(index):
with open(index, encoding='utf-8') as f:
self.weight_map = json.load(f)['weight_map']
else:
single = os.path.join(model_dir, 'diffusion_pytorch_model.safetensors')
with safe_open(single, framework='pt', device='cpu') as f:
self.weight_map = dict.fromkeys(f.keys(), 'diffusion_pytorch_model.safetensors')
def get(self, key):
shard = self.weight_map.get(key)
if shard is None:
return None
f = self.handles.get(shard)
if f is None:
f = safe_open(os.path.join(self.model_dir, shard), framework='pt', device='cpu')
self.handles[shard] = f
return f.get_tensor(key)
def quant_noise_energy(step, shape, hadamard_group, rms):
"""Expected energy of the base weight's quantization error, ``step^2/12`` per element under uniform rounding;
with ``rms`` per input channel it is the output-space energy per token."""
_out, n_in = shape
s = step.to(device, torch.float32)
if s.ndim == 3:
s = s.squeeze(-1)
if s.ndim == 1:
s = s[:, None]
groups = s.shape[1]
glen = n_in // groups
if rms is None:
power = torch.full((groups,), float(glen), device=s.device)
else:
p = rms.to(s.device, torch.float32).square()
if hadamard_group and n_in % hadamard_group == 0:
p = p.view(-1, hadamard_group).mean(1, keepdim=True).expand(-1, hadamard_group).reshape(-1)
power = p.view(groups, glen).sum(1)
return float(((s.square() / 12) @ power).sum())
def analyze_module(W_dq, deq_params, mods, calib_rms=None, step_live=None, noise=None):
"""Return fidelity metrics for one quantized module and the adapters targeting it.
Deltas come from each module's production calc_updown and sum the way the
loader stacks them, so every family (and dora / dense-bias / diff_b variant)
is measured as applied. A module is factor-path eligible only when every
contribution is a plain additive lora. With ``calib_rms``, hosting mirrors
the calibrated production path and its rho is scored in the weighted norm.
With ``step_live`` (the layer's own pre-add scale), the production routing
rule applies: a delta fat against the grid whose truncation capture is low
reports the requantize path, the way the loader would route it.
``snr`` is the delta's energy over the energy of the base weight's own
quantization error, in the calibrated norm when ``calib_rms`` is given:
``noise`` carries that error's measured ``(plain, weighted)`` energies
against an unquantized reference, else the uniform-rounding estimate from
``step_live`` stands in. It is independent of the apply path and says
whether what the adapter adds stands above the error the checkpoint
already carries.
"""
D = None
for mod in mods:
d = mod.calc_updown(W_dq)[0].to(device, torch.float32).reshape(W_dq.shape)
D = d if D is None else D + d
nD = D.norm()
control = deq_params['weights_dtype'] == 'bf16' # unquantized reference: the delta just rounds into bf16
factor_eligible = (not control) and all(lora_sdnq.get_module_factors(m, device, torch.bfloat16) is not None for m in mods)
if float(nD) == 0.0: # an all-zero delta (some full-rank extractions carry empty .diff): retention is undefined, not erased
return dict(rank=getattr(mods[0], 'dim', None), rms_delta=0.0, rms_weight=float(W_dq.pow(2).mean().sqrt()),
step_ratio=None, crossers=None, requant_rho=None, requant_resid=None,
factor_eligible=factor_eligible, hosted=False, applied_rho=None, delta_energy=0.0,
snr=None, snr_plain=None, noise_rms=None, noise_source=None)
snr, snr_plain, noise_rms, noise_source = None, None, None, None
if (noise is not None or step_live is not None) and not deq_params.get('use_codebook', False):
rms = calib_rms.to(device, torch.float32) if (calib_rms is not None and calib_rms.shape[-1] == D.shape[-1]) else None
if noise is not None:
plain, weighted = noise
noise_source = 'measured'
else:
hadamard = deq_params['hadamard_group_size'] if deq_params['use_hadamard'] else 0
plain = quant_noise_energy(step_live, D.shape, hadamard, None)
weighted = quant_noise_energy(step_live, D.shape, hadamard, rms) if rms is not None else plain
noise_source = 'uniform'
delta_w = float((D * rms).square().sum()) if rms is not None else float(nD.square())
noise_rms = (plain / D.numel()) ** 0.5
snr_plain = float(nD.square()) / plain if plain > 0 else None
snr = delta_w / weighted if weighted > 0 else None
step_ratio, crossers = None, None
if control:
W2 = (W_dq + D).to(torch.bfloat16).float()
else:
# mirror network_add_weights: it requantizes with the layer's own svd setting and rank,
# and an svd checkpoint's dequantized weight is not on the plain integer grid
use_svd = deq_params.get('use_svd', False)
kw = dict(layer_class_name='Linear', torch_dtype=torch.bfloat16, group_size=deq_params['group_size'],
hadamard_group_size=deq_params['hadamard_group_size'], use_hadamard=deq_params['use_hadamard'],
weights_dtype=deq_params['weights_dtype'], use_svd=use_svd, svd_rank=deq_params.get('svd_rank', 32),
svd_steps=deq_params.get('svd_steps', 8), use_quantized_matmul=False, dequantize_fp32=False)
deq2, data2 = sdnq_quantize_layer_weight(W_dq + D, **kw)
W2 = deq2(data2['weight'], data2['scale'], zero_point=data2['zero_point'],
svd_up=data2['svd_up'], svd_down=data2['svd_down'], dtype=torch.float32, skip_compile=True)
Dh = rotate_hadamard(D, group_size=deq_params['hadamard_group_size']) if deq_params['use_hadamard'] else D
step = data2['scale'].float()
Dg = Dh.unflatten(-1, (step.shape[1], -1)) if step.ndim == 3 else Dh
step_ratio = float((Dg.abs() / step).mean())
crossers = float((Dg.abs() > step / 2).float().mean())
E = W2 - W_dq
rho = float(E.flatten() @ D.flatten() / nD.square())
resid = float((E - D).norm() / nD)
hosted = False
if factor_eligible:
# the factor path stores the delta losslessly, but the dequantizer materializes
# base + factors in the result dtype (bf16 here), so realized fidelity floors at
# the same ULP rounding an unquantized bf16 model applies to a merged delta
base16 = W_dq.to(torch.bfloat16).float()
realized = (W_dq.to(torch.bfloat16) + D.to(torch.bfloat16)).float() - base16
applied_rho = float(realized.flatten() @ D.flatten() / nD.square())
else:
applied_rho = rho
if (not control) and cli_args.host_rank > 0:
from sdnq.common import dtype_dict
if dtype_dict[deq_params['weights_dtype']]['num_bits'] < 8:
# mirror lora_sdnq.apply_hosted: seeded svd truncation, realized through the bf16 materialize
q = min(cli_args.host_rank, *D.shape)
rms = None
if calib_rms is not None and calib_rms.shape[-1] == D.shape[-1]:
rms = calib_rms.to(D.device, torch.float32).clamp(min=1e-8)
Dw = D * rms if rms is not None else D
with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []):
torch.manual_seed(0)
U, S, V = torch.svd_lowrank(Dw, q=min(q + 64, *D.shape), niter=8)
energy = float(S[:q].square().sum() / Dw.square().sum().clamp(min=1e-30))
routed = False
if step_live is not None and not deq_params.get('use_svd', False):
sr = float(D.square().mean().sqrt() / step_live.float().mean())
routed = sr > lora_sdnq.REQUANT_RATIO and energy < lora_sdnq.REQUANT_ENERGY
if not routed: # the loader routes fat, genuinely-truncated deltas back to requantize
Dk = (U[:, :q] * S[:q]) @ V[:, :q].t()
if rms is not None:
Dk = Dk / rms
base16 = W_dq.to(torch.bfloat16).float()
realized = (W_dq.to(torch.bfloat16) + Dk.to(torch.bfloat16)).float() - base16
if rms is not None: # weighted norm: the diagonal-covariance output-error proxy the calibrated truncation optimizes
Dr = D * rms
applied_rho = float((realized * rms).flatten() @ Dr.flatten() / Dr.square().sum())
else:
applied_rho = float(realized.flatten() @ D.flatten() / nD.square())
hosted = True
return dict(rank=getattr(mods[0], 'dim', None), rms_delta=float(D.pow(2).mean().sqrt()), rms_weight=float(W_dq.pow(2).mean().sqrt()),
step_ratio=step_ratio, crossers=crossers, requant_rho=rho, requant_resid=resid,
factor_eligible=factor_eligible, hosted=hosted, applied_rho=applied_rho,
delta_energy=float(nD.square()), snr=snr, snr_plain=snr_plain, noise_rms=noise_rms, noise_source=noise_source)
def main():
args = cli_args
model_dir = resolve_model_dir(args.model)
arch_mod = resolve_arch(args.arch)
with open(os.path.join(model_dir, 'config.json'), encoding='utf-8') as f:
model_config = json.load(f)
pre_quantized = model_config.get('quantization_config') is not None
repo, bf16_repo, reference = None, None, None
quant_stamps, bf16_stamps, ref_stamps = {}, {}, {}
adapt = getattr(arch_mod, 'adapt_weights', None) # an arch refits deltas onto a live layout that differs from the trained one
if pre_quantized:
rprint(f'model: "{model_dir}" pre-quantized={pre_quantized}')
repo = QuantRepo(model_dir, model_config, arch=args.arch)
quant_stamps = stamp_index(repo.layers)
if args.reference:
reference = Bf16Repo(resolve_model_dir(args.reference))
ref_stamps = stamp_index(k[:-len('.weight')] for k in reference.weight_map if k.endswith('.weight'))
rprint(f'reference: "{reference.model_dir}" tensors={len(reference.weight_map)}')
else:
bf16_repo = Bf16Repo(model_dir)
bf16_stamps = stamp_index(k[:-len('.weight')] for k in bf16_repo.weight_map if k.endswith('.weight'))
if args.dtype is None:
rprint('model is not quantized and no --dtype given: loras apply exactly, nothing to analyze')
return 0
rprint(f'model: "{model_dir}" simulating dtype={args.dtype} group={args.group} hadamard={args.hadamard_group}')
calib_stats = {}
if args.calib:
with safe_open(os.path.expanduser(args.calib), framework='pt', device='cpu') as f:
calib_stats = {k: f.get_tensor(k) for k in f.keys()}
rprint(f'calib: "{args.calib}" layers={len(calib_stats)}')
noise_energies = {} # per module, shared by every lora: the base error does not depend on the adapter
def measure_noise(lname, err):
rms = calib_stats.get(lname)
plain = float(err.square().sum())
weighted = float((err * rms.to(device, torch.float32)).square().sum()) if rms is not None and rms.shape[-1] == err.shape[-1] else plain
noise_energies[lname] = (plain, weighted)
return noise_energies[lname]
report = {'model': model_dir, 'pre_quantized': pre_quantized, 'loras': []}
worst_effective = 1.0
def write_report():
if args.json:
with open(args.json, 'w', encoding='utf-8') as f:
json.dump(report, f, indent=2)
for lora_path in args.lora:
lora_path = os.path.expanduser(lora_path)
try:
mapped, census, chunked = map_lora_modules(lora_path, arch_mod)
net = network.Network(os.path.basename(lora_path), StubOnDisk(lora_path))
rows, unquantized, unmatched, failed, non_matrix = [], [], [], [], []
keys = sorted(mapped)
if not args.full and len(keys) > args.sample:
keys = keys[::max(1, len(keys) // args.sample)][:args.sample]
for path in keys:
entries = mapped[path]
lname = path
noise = None
if pre_quantized:
if path not in repo.layers:
lname = quant_stamps.get(path.replace('.', '_'), '')
layer = repo.layers.get(lname)
if layer is None:
unmatched.append(path)
continue
deq = getattr(layer, 'sdnq_dequantizer', None)
if deq is None:
unquantized.append(path)
continue
if len(deq.original_shape) != 2:
non_matrix.append(path)
continue
layer = repo.materialize(lname)
W_dq = deq(layer.weight, layer.scale, zero_point=layer.zero_point, svd_up=layer.svd_up, svd_down=layer.svd_down,
skip_quantized_matmul=deq.use_quantized_matmul, dtype=torch.float32, skip_compile=True).to(device)
params = dict(weights_dtype=deq.weights_dtype, group_size=deq.group_size, hadamard_group_size=deq.hadamard_group_size,
use_hadamard=deq.use_hadamard, use_svd=layer.svd_up is not None, svd_rank=deq.svd_rank, svd_steps=deq.svd_steps,
use_codebook=getattr(deq, 'use_codebook', False))
step_live = layer.scale.detach().to(device)
repo.release(lname)
sd_module = layer
if reference is not None:
noise = noise_energies.get(lname)
if noise is None:
W_ref = reference.get(f'{lname}.weight')
if W_ref is None:
W_ref = reference.get(f'{ref_stamps.get(lname.replace(".", "_"), "")}.weight')
if W_ref is not None and tuple(W_ref.shape) == tuple(W_dq.shape):
noise = measure_noise(lname, W_dq - W_ref.to(device, torch.float32))
else:
W = bf16_repo.get(f'{path}.weight')
if W is None:
W = bf16_repo.get(f'{bf16_stamps.get(path.replace(".", "_"), "")}.weight')
if W is None:
unmatched.append(path)
continue
if W.ndim != 2: # norm/scale targets (e.g. adaLN_modulation) are 1-D; the quantizer and the stub both expect a matrix
non_matrix.append(path)
continue
if args.dtype == 'bf16':
W_dq = W.to(device, torch.bfloat16).float()
params = dict(weights_dtype='bf16', group_size=0, hadamard_group_size=0, use_hadamard=False)
step_live = None
else:
deq0, data0 = sdnq_quantize_layer_weight(W.to(device, torch.float32), layer_class_name='Linear', weights_dtype=args.dtype,
group_size=args.group, hadamard_group_size=args.hadamard_group, use_hadamard=args.hadamard_group > 0,
use_svd=False, use_quantized_matmul=False, dequantize_fp32=False, torch_dtype=torch.bfloat16)
W_dq = deq0(data0['weight'], data0['scale'], zero_point=data0['zero_point'], svd_up=None, svd_down=None, dtype=torch.float32, skip_compile=True)
params = dict(weights_dtype=args.dtype, group_size=deq0.group_size, hadamard_group_size=deq0.hadamard_group_size, use_hadamard=deq0.use_hadamard)
step_live = data0['scale'].detach()
noise = noise_energies.get(path) or measure_noise(path, W_dq - W.to(device, torch.float32)) # the simulated grid's own error, measured
sd_module = make_stub(W.shape)
try:
if adapt is not None and pre_quantized:
entries = [(fam, (adapt(sd_module, path, w, transformers=[repo.model]) or w) if fam == 'lora' else w) for fam, w in entries]
mods = [build_module(fam, path, w, net, sd_module) for fam, w in entries]
row = analyze_module(W_dq, params, mods, calib_rms=calib_stats.get(lname), step_live=step_live, noise=noise)
except Exception as e: # a family the tool cannot rebuild must not read as a clean module
failed.append(f'{path}: {type(e).__name__}: {e}')
del W_dq
continue
row.update(module=path, dtype=params['weights_dtype'], family='+'.join(f for f, _w in entries))
rows.append(row)
del W_dq # the caching allocator reuses these; emptying it per module costs more than it saves
scored = [r for r in rows if r['applied_rho'] is not None] # zero-delta modules have no retention to report
applied = sorted(r['applied_rho'] for r in scored)
median_applied = applied[len(applied) // 2] if applied else None
energy = sum(r['delta_energy'] for r in scored)
weighted = (sum(r['applied_rho'] * r['delta_energy'] for r in scored) / energy) if energy > 0 else None
n_exact = sum(1 for r in scored if r['factor_eligible'])
fb = [r['requant_rho'] for r in scored if not r['factor_eligible']]
fb_median = sorted(fb)[len(fb) // 2] if fb else None
if median_applied is not None:
worst_effective = min(worst_effective, median_applied)
snr_rows = [r for r in scored if r['snr'] is not None]
snr_median = sorted(r['snr'] for r in snr_rows)[len(snr_rows) // 2] if snr_rows else None
snr_energy = sum(r['delta_energy'] for r in snr_rows)
snr_weighted = (sum(r['snr'] * r['delta_energy'] for r in snr_rows) / snr_energy) if snr_energy > 0 else None
report['loras'].append({'file': lora_path, 'families': census, 'targets': len(mapped), 'unquantized': unquantized,
'unmatched': unmatched, 'non_matrix': non_matrix, 'chunked': chunked, 'failed': failed,
'exact_modules': n_exact, 'fallback_modules': len(fb), 'fallback_median_rho': fb_median,
'median_applied_rho': median_applied, 'weighted_applied_rho': weighted,
'snr_median': snr_median, 'snr_weighted': snr_weighted, 'modules': rows})
write_report() # rewrite per file so a crash keeps completed work
rprint(f'\nlora: "{os.path.basename(lora_path)}" families={census or "none"} targets={len(mapped)} analyzed={len(rows)} scored={len(scored)} exact={n_exact} fallback={len(fb)} unquantized={len(unquantized)} unmatched={len(unmatched)} non_matrix={len(non_matrix)} chunked={chunked} failed={len(failed)}')
if median_applied is None:
rprint(' no analyzable modules: nothing measured')
else:
rprint(f' applied fidelity: median={median_applied:.3f} energy-weighted={weighted:.3f}' + (f' (fallback modules land at median rho={fb_median:.3f})' if fb_median is not None else ''))
if snr_median is not None:
rprint(f' delta over base quantization error: snr median={snr_median:.3g} energy-weighted={snr_weighted:.3g} ({snr_rows[0]["noise_source"]} noise, {"calibrated" if calib_stats else "weight-space"} norm)')
for f in failed[:3]:
rprint(f' [red]could not rebuild[/red]: {f}')
if fb:
worst = sorted((r for r in scored if not r['factor_eligible']), key=lambda r: r['requant_rho'])[:5]
rprint(' lowest-retention modules:')
for r in worst:
grid = f'step-ratio={r["step_ratio"]:.3f} crossers={r["crossers"]*100:5.1f}%' if r['step_ratio'] is not None else 'unquantized reference'
rprint(f' {r["module"]:48s} fam={r["family"]:5s} dtype={r["dtype"]} {grid} rho={r["requant_rho"]:.3f}')
del mapped, net
except KeyboardInterrupt:
raise
except Exception as e: # one broken file must not cost the rest of the batch
rprint(f'\n[red]lora failed[/red]: "{os.path.basename(lora_path)}" {type(e).__name__}: {e}')
report['loras'].append({'file': lora_path, 'error': f'{type(e).__name__}: {e}'})
write_report()
if device.type == 'cuda':
torch.cuda.empty_cache() # once per file, after its modules are done
report['complete'] = True
write_report()
if args.json:
rprint(f'\nreport: "{args.json}"')
if args.fail_under is not None and worst_effective < args.fail_under:
rprint(f'FAIL: effective fidelity {worst_effective:.3f} < {args.fail_under}')
return 2
return 0
if __name__ == '__main__':
with torch.inference_mode():
sys.exit(main())
cli/sdnq-attention-benchmark.py
+346 -56
View File
@@ -96,25 +96,75 @@ def save_transcript(path):
file_console.print(renderable)
console.print(f"results saved to {path}")
def key_padding_mask(cfg, device, keep=0.75):
# boolean key-padding mask over the kv axis, first keep fraction of keys valid
kv_tokens = cfg.get("kv_tokens", cfg["tokens"])
attn_mask = torch.zeros(cfg["batch"], 1, 1, kv_tokens, device=device, dtype=torch.bool)
attn_mask[..., :int(kv_tokens * keep)] = True
return attn_mask
def krea2_segment_mask(cfg, device):
# the transformer's segment_mask: text is padded to a fixed 512 tokens ahead of the
# image tokens and the padded tail is masked for queries and keys both, so padding
# query rows are fully masked and yield nan under sdpa (the model nan_to_num's them)
valid = torch.ones(cfg["batch"], cfg["tokens"], device=device, dtype=torch.bool)
valid[:, 128:512] = False
return valid.unsqueeze(1).unsqueeze(2) * valid.unsqueeze(1).unsqueeze(3)
def build_preset_mask(cfg, device):
# dense masks come from the preset's mask_fn; the element guard keeps h3-scale presets
# from materializing multi-gigabyte masks, those shapes belong to block-granular masks
mask_fn = cfg.get("mask_fn")
if mask_fn is None:
return None
attn_mask = mask_fn(cfg, device)
if attn_mask is not None and attn_mask.numel() > 2**31:
raise ValueError(f"preset dense mask holds {attn_mask.numel():,} elements; this shape needs a block-granular mask, not a token mask")
return attn_mask
shape_presets = {
# geometry from the model transformer and text-encoder configs;
# optional keys: kv_tokens (cross-attention), kv_heads (gqa), causal
# optional keys: kv_tokens (cross-attention), kv_heads (gqa), causal,
# mask_fn (token-granular attn_mask builder), mask_nan_guard (fully-masked query
# rows nan under stock sdpa), iters/warmup (per-preset run overrides for very large
# shapes), ref_head_chunk (head-sliced fp32 reference to bound peak memory),
# sparse (token layout driving the sparse selector rows)
"sd15": dict(batch=2, heads=8, tokens=4096, head_dim=40, desc="SD 1.5 unet self-attention at 512px, batched cfg, head dim padded 40 to 64"),
"sdxl": dict(batch=2, heads=10, tokens=4096, head_dim=64, desc="SDXL unet self-attention at 1024px, batched cfg"),
"sdxl-cross": dict(batch=2, heads=10, tokens=4096, kv_tokens=77, head_dim=64, desc="SDXL unet cross-attention at 1024px, 77 text tokens"),
"qwen3-te": dict(batch=2, heads=16, kv_heads=8, tokens=512, head_dim=128, causal=True, desc="Qwen3 text encoder (Anima), causal gqa 16:8 heads, 512 token prompt"),
"anima": dict(batch=1, heads=16, tokens=4096, head_dim=128, desc="Anima 1.0 self-attention at 1024px, one cfg pass"),
"flux2": dict(batch=1, heads=32, tokens=4608, head_dim=128, desc="FLUX.2 Klein 9B joint attention at 1024px, 4096 image plus 512 text tokens"),
"krea2": dict(batch=1, heads=48, tokens=4608, head_dim=128, desc="Krea 2 12B joint attention at 1024px, 4096 image plus 512 text tokens (128 real), kv expanded from gqa 48:12, segment mask"),
"krea2": dict(batch=1, heads=48, tokens=4608, head_dim=128, mask_fn=krea2_segment_mask, mask_nan_guard=True, desc="Krea 2 12B joint attention at 1024px, 4096 image plus 512 text tokens (128 real), kv expanded from gqa 48:12, segment mask"),
"wan22": dict(batch=1, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, one cfg pass"),
"wan22-cfg": dict(batch=2, heads=40, tokens=32760, head_dim=128, desc="Wan 2.2 A14B self-attention, 832x480 81 frames, batched cfg"),
"ltx2": dict(batch=1, heads=32, tokens=13376, head_dim=128, desc="LTX 2.3 self-attention, 1216x704 121 frames, one cfg pass"),
"masked": dict(batch=1, heads=32, tokens=4608, head_dim=128, desc="FLUX.2 Klein shape with boolean key-padding mask, 25% of keys masked"),
"h3": dict(batch=1, heads=56, tokens=38222, head_dim=128, iters=8, warmup=3, ref_head_chunk=14,
sparse=dict(layout=[("text", 0, 512), ("audio", 512, 926), ("video", 926, 38222)]),
desc="MiniMax H3 packed self-attention, 1344x768 124 frames (5.2s): 512 text + 414 audio + 37296 video rows, guidance-free"),
"h3-long": dict(batch=1, heads=56, tokens=109574, head_dim=128, iters=6, warmup=2, ref_head_chunk=8, config_timeout=1200,
sparse=dict(layout=[("text", 0, 512), ("audio", 512, 1718), ("video", 1718, 109574)]),
desc="MiniMax H3 packed self-attention, 1344x768 362 frames (15.1s): 512 text + 1206 audio + 107856 video rows"),
"masked": dict(batch=1, heads=32, tokens=4608, head_dim=128, mask_fn=key_padding_mask, desc="FLUX.2 Klein shape with boolean key-padding mask, 25% of keys masked"),
}
full_run = ["sd15", "sdxl", "sdxl-cross", "qwen3-te", "anima", "flux2", "krea2", "wan22", "ltx2"]
# sparse crossover probes at fixed h3 geometry; the smallest token count where a sparse row
# beats dense past the verdict threshold is the measured minimum-sequence gate
for gate_tokens in (2048, 4096, 8192, 16384, 32768, 65536):
shape_presets[f"gate-{gate_tokens // 1024}k"] = dict(
batch=1, heads=56, tokens=gate_tokens, head_dim=128,
sparse=dict(layout=[("text", 0, 512), ("video", 512, gate_tokens)]),
desc=f"sparse crossover probe at h3 geometry, {gate_tokens} tokens",
**(dict(iters=8, warmup=3) if gate_tokens >= 32768 else {}),
)
full_run = ["sd15", "sdxl", "sdxl-cross", "qwen3-te", "anima", "flux2", "krea2", "wan22", "ltx2", "h3"]
sparse_run = ["krea2", "h3", "h3-long"]
gate_run = [f"gate-{tokens // 1024}k" for tokens in (2048, 4096, 8192, 16384, 32768, 65536)]
# settings advice comes from a self-attention shape with the full config set; cross-attention
# and text-encoder shapes measure the hijack's cost there but would mislead as global advice
recommendation_presets = ["flux2", "krea2", "anima", "sdxl", "wan22", "ltx2", "sd15"]
recommendation_presets = ["flux2", "krea2", "anima", "sdxl", "wan22", "ltx2", "h3", "sd15"]
default_shapes = "sdxl,flux2"
all_sections = ["attention", "dequant", "block"]
@@ -122,11 +172,16 @@ all_sections = ["attention", "dequant", "block"]
# measures a complete configuration of weights dtype x matmul path x attention end to end
# generic dit-block geometries from the model transformer configs: flux.1 (3072 wide,
# 24 heads, 4x gelu ff, 4096 image plus 512 text tokens) and krea 2 (6144 wide, 48 heads
# after gqa expansion, swiglu at 16384, same joint sequence)
# after gqa expansion, swiglu at 16384, same joint sequence); optional keys: head_dim
# (attention width when heads*head_dim != hidden) and mlp ("gelu" default or "swiglu")
block_geometries = {
"flux1": dict(hidden=3072, heads=24, mlp_dim=12288, tokens=4608),
"krea2": dict(hidden=6144, heads=48, mlp_dim=16384, tokens=4608),
# minimax h3: attention wider than the residual stream (56*128 > 5376), swiglu mlp; the
# full 124-frame token count makes the block section long, so it runs only when selected
"h3": dict(hidden=5376, heads=56, head_dim=128, mlp_dim=14336, mlp="swiglu", tokens=38222, iters=6, warmup=2, config_timeout=900),
}
default_block_geometries = "flux1,krea2"
block_geometry = block_geometries["flux1"] # active geometry; bench_block_section iterates
block_attention_specs = {
"sdpa": None, # stock torch sdpa
@@ -139,6 +194,19 @@ block_attention_specs = {
"sage": "sage", # external baselines, resolved to the sage wrappers in build_bench_block
"sage fp16 accum": "sagefp16",
}
# attention-table config id measuring the same kernel as each block attention spec, for the
# cross-instrument compute split; specs without a standalone row map to None
block_spec_attention_ids = {
"sdpa": "base",
"atten int8": "int8",
"atten int8 smooth": "smooth",
"atten int8 hadamard": "hadamard",
"atten full": "full",
"atten pv accum": "pvaccum",
"atten fp16 accum": "fp16full-accum",
"sage": "sage",
"sage fp16 accum": "sagefp16",
}
# id, weights config (None = bf16), use quantized matmul, attention spec; fp8/fp4 rows use the
# dequant path: quantized matmul auto-selects fp8 for float dtypes, unsupported before sm_89
block_configs = [
@@ -173,8 +241,19 @@ bench_configs = [
("sage", "sageattention", None), # label resolved to the dispatched kernel by sage_kernel_label
("sagefp16", "sage int8 qk + fp16 pv, fp16 accum", None), # sm86 only
("amdflash", "triton flash (amd)", None),
("flex", "flex attention, dense", None), # compiled: flex reads its block lists only under compile
("flex-sparse100", "flex + selector, budget 100%", None), # the selector runs but keeps everything, so this row is its overhead alone
("flex-sparse50", "flex + selector, budget 50%", None),
("flex-sparse30", "flex + selector, budget 30%", None),
("flex-sparse15", "flex + selector, budget 15%", None),
("flex-radial30", "flex + static radial band, 30%", None), # density matched control with no per-call producer
("noquant", "sdnq, quantized matmul off", dict(do_quantize=False)),
("int8", "sdnq int8 qk", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
("int8-sparse100", "sdnq int8 qk + selector, budget 100%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # the selector runs but keeps everything, so this row is its overhead on the quantized kernel
("int8-sparse50", "sdnq int8 qk + selector, budget 50%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
("int8-sparse30", "sdnq int8 qk + selector, budget 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
("int8-sparse15", "sdnq int8 qk + selector, budget 15%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")),
("int8-radial30", "sdnq int8 qk + static radial band, 30%", dict(matmul_dtype="auto", pv_matmul_dtype="auto")), # density matched control with no per-call producer
("smooth", "sdnq int8 qk + smooth k", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True)),
("hadamard", "sdnq int8 qk + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", use_hadamard=True)),
("smooth_hadamard", "sdnq int8 qk + smooth + hadamard", dict(matmul_dtype="auto", pv_matmul_dtype="auto", smooth_k=True, use_hadamard=True)),
@@ -193,7 +272,10 @@ bench_configs = [
# external baselines are compared against but never starred or recommended as sdnq configs;
# the unsafe accum mode is measured and displayed under the same rule, since its overflow
# tail lives outside what mean error can see
external_config_ids = ("base", "sage", "sagefp16", "amdflash")
# baselines and non-sdnq rows: reported, never recommended as an sdnq setting, and the sparse
# rows are lossy by design so a recommendation must not pick one for being fast
external_config_ids = ("base", "sage", "sagefp16", "amdflash", "flex")
sparse_config_ids = ("flex-sparse100", "flex-sparse50", "flex-sparse30", "flex-sparse15", "flex-radial30", "int8-sparse100", "int8-sparse50", "int8-sparse30", "int8-sparse15", "int8-radial30")
unsafe_config_ids = ("pvaccum",)
# every preset runs the full config list (availability gates still apply per config); only
# hard technical exclusions live here, never runtime trims. sd15: compiling hadamard with
@@ -281,18 +363,22 @@ def parse_cli():
parser.add_argument("--dequant-sweeps", type=str, default="all", help=f"comma-separated setting sweeps in the dequant section: {', '.join(all_dequant_sweeps)}; 'all' or 'none' (default: %(default)s)")
parser.add_argument("--mm-backends", type=str, default="none", help=f"comma-separated quantized-matmul backends to compare in one run: {', '.join(all_mm_backends)}; 'none' benches only the backend this device selects (default: %(default)s)")
parser.add_argument("--mm-rounds", type=int, default=2, help="alternating rounds per matmul backend, fastest kept, so clock drift cancels instead of favouring one backend (default: %(default)s)")
parser.add_argument("--configs", type=str, default="all", help=f"comma-separated attention configs: {', '.join(config_id for config_id, _label, _kwargs in bench_configs)}; 'all' runs every one (default: %(default)s)")
parser.add_argument("--block-configs", type=str, default="all", help=f"comma-separated combined block configs: {', '.join(config_id for config_id, _w, _mm, _a in block_configs)}; 'all' runs every one (default: %(default)s)")
parser.add_argument("--shapes", type=str, default=default_shapes, help=f"comma-separated attention shape presets: {', '.join(shape_presets)}; 'all' runs {', '.join(full_run)} (default: %(default)s)")
parser.add_argument("--block-geometries", type=str, default=default_block_geometries, help=f"comma-separated block geometries: {', '.join(block_geometries)}; 'all' runs every one (default: %(default)s)")
parser.add_argument("--shapes", type=str, default=default_shapes, help=f"comma-separated attention shape presets: {', '.join(shape_presets)}; 'all' runs {', '.join(full_run)}, 'sparse' and 'gate' run the sparse and crossover lists (default: %(default)s)")
parser.add_argument("--iters", type=int, default=12, help="minimum timed iterations per config, scaled up for fast kernels (default: %(default)s)")
parser.add_argument("--warmup", type=int, default=4, help="minimum warmup iterations per config, scaled up for fast kernels (default: %(default)s)")
parser.add_argument("--skip-checks", action="store_true", help="skip kernel correctness checks")
parser.add_argument("--skip-bench", action="store_true", help="skip benchmarks, run checks and the fp8 and compile probes only")
parser.add_argument("--dtype", type=str, default="auto", choices=["auto", "bf16", "fp16"], help="tensor dtype for benchmarks; auto uses the dtype the webui selected for this gpu (default: %(default)s)")
parser.add_argument("--config-timeout", type=int, default=300, help="best effort: abort a config whose compile plus first call exceeds this many seconds, 0 disables; cannot interrupt native-level hangs (default: %(default)s)")
parser.add_argument("--config-timeout", type=int, default=None, help="best effort: abort a config whose compile plus first call exceeds this many seconds, 0 disables; cannot interrupt native-level hangs (default: 300, or the limit a preset or block geometry declares for itself)")
parser.add_argument("--save", type=str, default="auto", help="plain-text copy of all tables and notes; 'auto' (default) names it <gpu>-t<torch>-<date>.txt in the output directory, 'none' disables, anything else is used as the path")
parser.add_argument("--json", type=str, default="auto", help="structured results (environment, probes, per-shape and dequant timings, recommendations); 'auto' (default) names it <gpu>-t<torch>-<date>.json in the output directory, 'none' disables, anything else is used as the path")
parser.add_argument("--outdir", type=str, default=None, help="directory for auto-named outputs (default: $SDNQ_BENCH_DIR, or benchmarks/ under the sdnext root)")
args = parser.parse_args()
args.timeout_flag = args.config_timeout # None lets a preset or block geometry declare its own limit
args.config_timeout = resolve_timeout(args.timeout_flag)
sys.argv = sys.argv[:1] # sdnext parses argv again on import and rejects unknown arguments
return args
@@ -379,7 +465,7 @@ def load_sdnext():
raise
detail = f"exited with code {e.code}" if isinstance(e, SystemExit) else f"{type(e).__name__}: {e}"
console.print(f"[red]sdnext failed to start: {detail}[/red]")
text = startup_log["text"].strip()
text = startup_log["text"].strip() # pylint: disable=used-before-assignment
if text:
console.print(Panel(escape(text[-4000:]), title="sdnext startup log", box=ROUNDED_BOX))
console.print("run from the sdnext root with the venv active; triton is required")
@@ -425,6 +511,16 @@ def atten_supports_fp16_accum():
return False
def atten_supports_block_mask():
# the sdnq sparse rows feed the kernel's block_mask kwarg; skip them on builds without it
if sdnq_triton_atten is None:
return False
try:
return "block_mask" in inspect.signature(inspect.unwrap(sdnq_triton_atten)).parameters
except (TypeError, ValueError):
return False
def triton_mm_supports_fp16_accum():
try:
from sdnq.kernels import triton_mm, triton_scaled_mm
@@ -450,6 +546,56 @@ def triton_mm_fp16_accum():
triton_mm.USE_FP16_ACCUM, triton_scaled_mm.USE_FP16_ACCUM = saved
flex_budgets = {"flex-sparse100": 1.0, "flex-sparse50": 0.50, "flex-sparse30": 0.30, "flex-sparse15": 0.15}
sdnq_sparse_budgets = {"int8-sparse100": 1.0, "int8-sparse50": 0.50, "int8-sparse30": 0.30, "int8-sparse15": 0.15}
def is_sdnq_sparse(config_id):
return config_id in sdnq_sparse_budgets or config_id == "int8-radial30"
def make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa):
"""The producer the flex rows time, feeding the quantized kernel's block mask input instead."""
from modules.attention.sparse import selector as sparse_selector
def attend(selection):
return sdnq_triton_atten(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, **kwargs)
if config_id == "int8-radial30":
static = sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device)
return lambda: attend(static)
spec = sparse_selector.SparseSpec(budget=sdnq_sparse_budgets[config_id], force=True)
cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape))
return lambda: attend(sparse_selector.select_blocks(q, k, spec, cache_key=cache_key))
def flex_available():
try:
import modules.attention.sparse.flex # pylint: disable=unused-import
return torch.cuda.is_available()
except Exception:
return False
def make_flex_fn(config_id, q, k, v, scale, gqa):
"""Time the selector inside the attention it accelerates; a producer measured on its own looks free and is not."""
from modules.attention.sparse import flex as sparse_flex, selector as sparse_selector
call = sparse_flex.flex_call()
if config_id == "flex":
return lambda: call(q, k, v, scale=scale, enable_gqa=gqa)
if config_id == "flex-radial30":
# a static pattern is built once by construction, which is exactly the advantage it has to defend
static = sparse_flex.to_block_mask(sparse_selector.radial_blocks(q.shape[-2], k.shape[-2], 0.30, sparse_selector.SparseSpec(), q.device))
return lambda: call(q, k, v, block_mask=static, scale=scale, enable_gqa=gqa)
spec = sparse_selector.SparseSpec(budget=flex_budgets[config_id], force=True)
cache_key = ("bench", config_id, tuple(q.shape), tuple(k.shape)) # the webui caches the geometry per layout, so measure that path
def run():
selection = sparse_selector.select_blocks(q, k, spec, cache_key=cache_key)
return call(q, k, v, block_mask=sparse_flex.to_block_mask(selection), scale=scale, enable_gqa=gqa)
return run
def sage_attention():
# mirror the backend selection from modules/attention.py: sm86 needs the cuda backend
try:
@@ -545,13 +691,23 @@ def make_qkv(batch, heads, tokens, head_dim, structured=True, kv_heads=None, kv_
return q, k, v
def fp32_reference(q, k, v, **kwargs):
def fp32_reference(q, k, v, head_chunk=0, **kwargs):
# sdnext enables tf32 globally; a math-backend dispatch fallback would degrade the reference to tf32 precision
tf32_matmul = torch.backends.cuda.matmul.allow_tf32
tf32_cudnn = torch.backends.cudnn.allow_tf32
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
try:
if head_chunk and not kwargs.get("enable_gqa") and q.shape[1] > head_chunk:
# head-sliced reference: bounds the fp32 peak on very long sequences; gqa shapes
# keep the one-shot path since slicing q heads would have to regroup kv heads
attn_mask = kwargs.pop("attn_mask", None)
outs = []
for start in range(0, q.shape[1], head_chunk):
heads = slice(start, start + head_chunk)
mask_slice = attn_mask[:, heads] if attn_mask is not None and attn_mask.shape[1] > 1 else attn_mask
outs.append(torch.nn.functional.scaled_dot_product_attention(q[:, heads].to(torch.float32), k[:, heads].to(torch.float32), v[:, heads].to(torch.float32), attn_mask=mask_slice, **kwargs))
return torch.cat(outs, dim=1)
return torch.nn.functional.scaled_dot_product_attention(q.to(torch.float32), k.to(torch.float32), v.to(torch.float32), **kwargs)
finally:
torch.backends.cuda.matmul.allow_tf32 = tf32_matmul
@@ -641,6 +797,13 @@ def live_progress():
return progress, task
def resolve_timeout(flag, declared=None):
# the cli flag wins when given; otherwise a preset or block geometry may declare its own limit
if flag is not None:
return flag
return 300 if declared is None else declared
@contextmanager
def time_limit(seconds, label):
# torch.compile can spin indefinitely in sympy/inductor on pathological graphs
@@ -721,6 +884,9 @@ def run_drift_sigma():
drift_override = None
verdict_z = 1.28 # one-sided 90%: an on/off verdict is only stated when its margin test clears this
# split instruments A and B differ systematically (strided views out of the fused projection vs
# contiguous standalone tensors): 0.1-2.5% over 19 same-kernel rows at h3, far below the 17-25% hadamard gap
split_instrument_offset = 0.05
def sidak_z_for(count):
@@ -826,7 +992,7 @@ def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_resul
lines.append(f"float8_e4m3fn matmul: [red]not supported on this gpu, selecting it fails generation[/red] [dim]({escape(fp8_result['qk'][1])})[/dim]")
else:
lines.append(f"float8_e4m3fn matmul: [red]failed to compile in this environment, selecting it fails generation[/red]; the error is not the hardware-capability signature, a torch or triton issue is more likely than the gpu [dim]({escape(fp8_result['qk'][1])})[/dim]")
lines.append(f"sdnq attention enabled in current config: {'[green]yes[/green]' if 'SDNQ attention' in shared.opts.sdp_overrides else '[yellow]no, enable via Compute Settings -> SDP overrides (requires restart)[/yellow]'}")
lines.append(f"sdnq attention enabled in current config: {'[green]yes[/green]' if 'SDNQ attention' in shared.opts.cross_attention_optimization else '[yellow]no, enable via Compute Settings -> Cross Attention (requires restart)[/yellow]'}")
if prep_status == "disabled":
lines.append("compiled input prep: torch.compile disabled in config, input prep runs eager")
elif prep_status == "working":
@@ -853,6 +1019,8 @@ def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_resul
lines.append(f"compiled weight dequant, float8_e5m2 storage: {e5m2_verdict}")
if not atten_supports_fp16_accum():
lines.append("fp16 accumulation kwarg: [yellow]absent in this sdnq build, accum rows skipped[/yellow]")
if not atten_supports_block_mask():
lines.append("block mask kwarg: [yellow]absent in this sdnq build, sdnq sparse rows skipped[/yellow]")
overrides = [f"{key}={value}" for key, value in os.environ.items() if key.startswith("SDNQ_TRITON_ATTEN") or key.startswith("SDNQ_TRITON_MM") or key.startswith("SDNQ_ALLOW_FP8") or key.startswith("SDNQ_COMPILE")]
if overrides:
lines.append(f"env overrides: {' '.join(overrides)}")
@@ -870,6 +1038,7 @@ def print_environment(fp8_result, prep_status, prep_detail, weight_dequant_resul
**runtime_versions,
fp8_attention_matmul=fp8_result["qk"][0] if fp8_result is not None else None,
atten_fp16_accum=atten_supports_fp16_accum(),
atten_block_mask=atten_supports_block_mask(),
triton_mm_fp16_accum=os.environ.get("SDNQ_TRITON_MM_USE_FP16_ACCUM", None),
compiled_input_prep=prep_status,
fp8_compile_gate=fp8_compile_gate_flag(),
@@ -1009,35 +1178,56 @@ def make_source_weight(out_features, in_features, seed=1234):
class BenchBlock(torch.nn.Module):
# dit-style block: fused qkv self-attention plus a gelu mlp, both with residuals; the
# attention_fn attribute is set per benchmark config (stock sdpa or sdnq attention)
def __init__(self, hidden, heads, mlp_dim, device=None, dtype=None):
# dit-style block: fused qkv self-attention plus a gelu or swiglu mlp, both with residuals;
# head_dim decouples attention width from hidden for models whose attention is wider than
# the residual stream; the attention_fn attribute is set per benchmark config
def __init__(self, hidden, heads, mlp_dim, head_dim=None, mlp="gelu", device=None, dtype=None):
super().__init__()
self.heads = heads
self.head_dim = head_dim if head_dim is not None else hidden // heads
self.mlp = mlp
inner = heads * self.head_dim
self.norm1 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype)
self.norm2 = torch.nn.LayerNorm(hidden, elementwise_affine=False, device=device, dtype=dtype)
self.qkv = torch.nn.Linear(hidden, hidden * 3, bias=False, device=device, dtype=dtype)
self.proj = torch.nn.Linear(hidden, hidden, bias=False, device=device, dtype=dtype)
self.qkv = torch.nn.Linear(hidden, inner * 3, bias=False, device=device, dtype=dtype)
self.proj = torch.nn.Linear(inner, hidden, bias=False, device=device, dtype=dtype)
self.up = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype)
if mlp == "swiglu":
self.gate = torch.nn.Linear(hidden, mlp_dim, bias=False, device=device, dtype=dtype)
self.down = torch.nn.Linear(mlp_dim, hidden, bias=False, device=device, dtype=dtype)
self.attention_fn = None
def forward(self, x):
batch, tokens, channels = x.shape
batch, tokens, _channels = x.shape
h = self.norm1(x)
qkv = self.qkv(h).view(batch, tokens, 3, self.heads, channels // self.heads).permute(2, 0, 3, 1, 4)
attn = self.attention_fn(qkv[0], qkv[1], qkv[2]).transpose(1, 2).reshape(batch, tokens, channels)
qkv = self.qkv(h).view(batch, tokens, 3, self.heads, self.head_dim).permute(2, 0, 3, 1, 4)
attn = self.attention_fn(qkv[0], qkv[1], qkv[2]).transpose(1, 2).reshape(batch, tokens, self.heads * self.head_dim)
x = x + self.proj(attn)
h = self.norm2(x)
if self.mlp == "swiglu":
return x + self.down(torch.nn.functional.silu(self.gate(h)) * self.up(h))
return x + self.down(torch.nn.functional.gelu(self.up(h)))
def build_block_module(dtype=None):
# construct a block for the active geometry; every construction site goes through here so
# geometry keys are read in exactly one place
return BenchBlock(
block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"],
head_dim=block_geometry.get("head_dim"), mlp=block_geometry.get("mlp", "gelu"),
device=torch_device, dtype=dtype if dtype is not None else bench_dtype,
)
def make_block_master():
# one master weight set shared by every block config, so all rows quantize identical weights
hidden, heads, mlp_dim = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"]
block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=bench_dtype)
# one master weight set shared by every block config, so all rows quantize identical weights;
# seed order keeps gelu geometries bitwise stable, swiglu appends its gate after up
block = build_block_module()
linears = [block.qkv, block.proj, block.up, block.down]
if hasattr(block, "gate"):
linears.append(block.gate)
with torch.no_grad():
for seed, linear in enumerate((block.qkv, block.proj, block.up, block.down), start=1):
for seed, linear in enumerate(linears, start=1):
linear.weight.copy_(make_source_weight(linear.out_features, linear.in_features, seed=seed))
return {key: value.clone() for key, value in block.state_dict().items()}
@@ -1045,8 +1235,7 @@ def make_block_master():
def build_bench_block(master_sd, weights_cfg, use_mm, attention_spec):
from sdnq import SDNQConfig
from sdnq.quantizer import apply_sdnq_to_module
hidden, heads, mlp_dim = block_geometry["hidden"], block_geometry["heads"], block_geometry["mlp_dim"]
block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=bench_dtype)
block = build_block_module()
block.load_state_dict(master_sd)
block.eval()
for param in block.parameters():
@@ -1335,36 +1524,29 @@ def make_prep_fn(q, k, v, attn_mask, kwargs, is_causal=False, enable_gqa=False):
return prep
def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_result=None):
def bench_shape(preset, iters, warmup, position=None, config_timeout=None, fp8_result=None, selected=None):
preset_cfg = shape_presets[preset]
config_timeout = resolve_timeout(config_timeout, preset_cfg.get("config_timeout"))
batch, heads, tokens, head_dim = preset_cfg["batch"], preset_cfg["heads"], preset_cfg["tokens"], preset_cfg["head_dim"]
kv_tokens = preset_cfg.get("kv_tokens", tokens)
kv_heads = preset_cfg.get("kv_heads", heads)
causal = preset_cfg.get("causal", False)
gqa = kv_heads != heads
description = preset_cfg["desc"]
iters = preset_cfg.get("iters", iters)
warmup = preset_cfg.get("warmup", warmup)
ref_head_chunk = preset_cfg.get("ref_head_chunk", 0)
excluded_configs = preset_excluded_configs.get(preset, set())
if preset == "sd15":
emit("[yellow]sd15: hadamard configs skipped, compiling hadamard with a non pow2 head dim currently hangs torch inductor[/yellow]")
attn_mask = None
mask_nan_guard = False
if preset == "masked":
attn_mask = torch.zeros(batch, 1, 1, tokens, device=torch_device, dtype=torch.bool)
attn_mask[..., :int(tokens * 0.75)] = True
elif preset == "krea2":
# the transformer's segment_mask: text is padded to a fixed 512 tokens ahead of the
# image tokens and the padded tail is masked for queries and keys both, so padding
# query rows are fully masked and yield nan under sdpa (the model nan_to_num's them)
valid = torch.ones(batch, tokens, device=torch_device, dtype=torch.bool)
valid[:, 128:512] = False
attn_mask = valid.unsqueeze(1).unsqueeze(2) * valid.unsqueeze(1).unsqueeze(3)
mask_nan_guard = True
attn_mask = build_preset_mask(preset_cfg, torch_device)
mask_nan_guard = preset_cfg.get("mask_nan_guard", False)
sage = sage_attention()
sage_fp16 = sage_attention_fp16_accum()
amd_flash = amd_triton_flash()
selected_configs = []
for config_id, label, kwargs in bench_configs:
if config_id in excluded_configs:
if config_id in excluded_configs or (selected is not None and config_id not in selected):
continue
if config_id == "sage" and (sage is None or attn_mask is not None or head_dim not in {64, 96, 128} or kv_tokens != tokens or gqa or causal):
continue
@@ -1372,6 +1554,10 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
continue
if config_id == "amdflash" and (amd_flash is None or attn_mask is not None or head_dim > 128 or gqa):
continue
if config_id.startswith("flex") and (not flex_available() or attn_mask is not None or causal or kv_tokens != tokens):
continue # a block only mask cannot carry a token mask or a causal rule, and cross attention is not sparsified
if is_sdnq_sparse(config_id) and (not atten_supports_block_mask() or causal or kv_tokens != tokens):
continue # the kernel composes a token mask with the block mask, so only the causal and cross attention rules apply
if config_id == "fp8qk" and not (fp8_result and fp8_result["qk"][0]):
continue
if config_id == "fp8pv" and not (fp8_result and fp8_result["pv"][0]):
@@ -1413,7 +1599,7 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
progress.update(task, description=f"{prefix}{preset}: preparing inputs and fp32 reference")
q, k, v = make_qkv(batch, heads, tokens, head_dim, kv_heads=kv_heads, kv_tokens=kv_tokens)
scale = head_dim ** -0.5
ref = fp32_reference(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa)
ref = fp32_reference(q, k, v, attn_mask=attn_mask, is_causal=causal, enable_gqa=gqa, head_chunk=ref_head_chunk)
if mask_nan_guard:
ref = torch.nan_to_num(ref)
anchor_fn = None
@@ -1432,6 +1618,10 @@ def bench_shape(preset, iters, warmup, position=None, config_timeout=300, fp8_re
elif config_id == "amdflash":
def fn(sm=scale):
return amd_flash(q, k, v, sm, is_causal=causal)
elif config_id.startswith("flex"):
fn = make_flex_fn(config_id, q, k, v, scale, gqa)
elif is_sdnq_sparse(config_id):
fn = make_sdnq_sparse_fn(config_id, q, k, v, attn_mask, kwargs, causal, gqa)
else:
def fn(kw=kwargs, mask=attn_mask):
return sdnq_triton_atten(q, k, v, attn_mask=mask, is_causal=causal, enable_gqa=gqa, **kw)
@@ -2039,7 +2229,7 @@ def resolved_group_label(layer, in_features):
return "row"
def bench_group_sizes(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300):
def bench_group_sizes(shape_label, out_features, in_features, plain_results, selected_dtypes, iters, warmup, config_timeout=300): # pylint: disable=unused-argument
# the Group size setting: 0 = auto, -1 = row-wise, explicit values snap to a divisor of
# in_features; grouping forces a per-forward re-quantize when quantized matmul is on, so
# the mm cells price that cost alongside the accuracy gain
@@ -2514,17 +2704,24 @@ def block_label(weights_cfg, use_mm, attention_spec):
return f"{weights_part} + {attention_spec}"
def bench_block_section(iters, warmup, config_timeout=300, selected=None):
def bench_block_section(iters, warmup, config_timeout=None, selected=None, geometries=None):
global block_geometry # pylint: disable=global-statement
all_results = {}
for family, geometry in block_geometries.items():
if geometries is not None and family not in geometries:
continue
block_geometry = geometry
results = bench_block_geometry(iters, warmup, config_timeout=config_timeout, selected=selected)
geometry_iters = geometry.get("iters", iters)
geometry_warmup = geometry.get("warmup", warmup)
geometry_timeout = resolve_timeout(config_timeout, geometry.get("config_timeout"))
results = bench_block_geometry(geometry_iters, geometry_warmup, config_timeout=geometry_timeout, selected=selected)
all_results[family] = results
report.setdefault("blocks", {})[family] = dict(geometry=dict(geometry), results=results)
# the first family also lands at the flat block key, which replays and the buyback
report.setdefault("blocks", {})[family] = dict(geometry=dict(geometry), results=results, split=None)
if not all_results:
return {}
# the first family run also lands at the flat block key, which replays and the buyback
# veto fall back to when no family matches the reference shape
primary = next(iter(block_geometries))
primary = next(iter(all_results))
report["block"] = report["blocks"][primary]
return all_results.get(primary, {})
@@ -2556,7 +2753,7 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
table.add_column("out err", justify="right")
table.add_column("max tok err", justify="right")
table.add_column("err x4 blocks", justify="right")
panel = Panel(table, title=f"combined block: hidden={hidden} heads={heads} mlp={mlp_dim} tokens={tokens} {dtype_label()}", subtitle="[dim]dit block, fused qkv + gelu mlp with residuals; err vs an fp32 reference block, max tok = worst single token, x4 = four stacked blocks[/dim]", box=ROUNDED_BOX, expand=False)
panel = Panel(table, title=f"combined block: hidden={hidden} heads={heads} mlp={mlp_dim} tokens={tokens} {dtype_label()}", subtitle="[dim]dit block, fused qkv attention + mlp with residuals; err vs an fp32 reference block, max tok = worst single token, x4 = four stacked blocks[/dim]", box=ROUNDED_BOX, expand=False)
def run_depth(block, x0, depth):
h = x0
@@ -2572,7 +2769,7 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
master = make_block_master()
generator = torch.Generator(device=torch_device).manual_seed(7)
x = torch.randn(1, tokens, hidden, device=torch_device, dtype=bench_dtype, generator=generator)
ref_block = BenchBlock(hidden, heads, mlp_dim, device=torch_device, dtype=torch.float32)
ref_block = build_block_module(dtype=torch.float32)
ref_block.load_state_dict(master)
ref_block.eval()
def ref_attention(q, k, v):
@@ -2619,6 +2816,15 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
phase("measuring depth-4 error")
with torch.no_grad():
entry["err4"] = rel_err(run_depth(block, x, 4), ref_out4)
phase("timing identity-attention variant")
real_attention_fn = block.attention_fn
def identity_attention_fn(q, k, v): # pylint: disable=unused-argument # same shapes and permutes, zero attention flops
return v
block.attention_fn = identity_attention_fn
try:
entry["identity_ms"], entry["identity_ms_sigma"] = bench_stats(fn, warmup, iters, on_phase=phase)
finally:
block.attention_fn = real_attention_fn
del block, out
if base_ms is None:
base_ms = entry["ms"]
@@ -2643,7 +2849,7 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
weights_mode = str(getattr(shared.opts, "sdnq_quantize_weights_mode", ""))
current_id = None
if weights_mode == "int8" and getattr(shared.opts, "sdnq_quantize_matmul_mode", "disabled") != "disabled":
current_id = "int8-mm-atten" if "SDNQ attention" in shared.opts.sdp_overrides else "int8-mm"
current_id = "int8-mm-atten" if "SDNQ attention" in shared.opts.cross_attention_optimization else "int8-mm"
if current_id and results.get(current_id, {}).get("ms"):
notes.append(f"current config runs the {results[current_id]['label']} row for int8-quantized models")
if any(entry.get("ms") for config_id, entry in results.items() if config_id.endswith("sagefp16")):
@@ -2653,6 +2859,72 @@ def bench_block_geometry(iters, warmup, config_timeout=300, selected=None):
return results
def emit_block_splits():
# instrument B reads the attention tables, so the split renders once both sections are in
for family, data in (report.get("blocks") or {}).items():
data["split"] = block_split_table(family, data["geometry"], data["results"])
def block_split_table(family, geometry, results):
# compute split per config from two independent instruments: A subtracts the identity-
# attention variant timed inside the block, B reads the standalone attention table at the
# same geometry from this run; a speedup ceiling is only stated where the two agree
head_dim = geometry.get("head_dim") or geometry["hidden"] // geometry["heads"]
expected_geometry = f"batch=1 heads={geometry['heads']} tokens={geometry['tokens']} head_dim={head_dim}"
attention_preset = None
for preset_name, data in (report.get("attention") or {}).items():
if data.get("geometry") == expected_geometry and not shape_presets.get(preset_name, {}).get("mask_fn"):
attention_preset = preset_name
break
attention_results = (report.get("attention") or {}).get(attention_preset, {}).get("results", {}) if attention_preset else {}
spec_by_config = {config_id: spec for config_id, _w, _mm, spec in block_configs}
budgets = (0.5, 0.3, 0.15)
table = Table(box=box.SIMPLE_HEAVY)
table.add_column("config")
table.add_column("block ms", justify="right")
table.add_column("rest ms", justify="right")
table.add_column("attn A", justify="right")
table.add_column("attn B", justify="right")
table.add_column("agree", justify="right")
for budget in budgets:
table.add_column(f"ceil@{int(budget * 100)}%", justify="right")
split = {}
for config_id, entry in results.items():
ms, identity_ms = entry.get("ms"), entry.get("identity_ms")
if not ms or not identity_ms:
continue
attn_a = ms - identity_ms
if attn_a <= 0:
continue
attention_id = block_spec_attention_ids.get(spec_by_config.get(config_id))
attention_entry = attention_results.get(attention_id) or {}
attn_b = attention_entry.get("ms")
agree = None
if attn_b:
# subtraction amplifies the relative sigma of instrument A by ms/attn_a
sigma_a = (row_sigma(entry) or 0.0) * (ms / attn_a)
sigma_b = row_sigma(attention_entry, "ms") or 0.0
threshold = max(verdict_z * math.sqrt(sigma_a * sigma_a + sigma_b * sigma_b), run_drift_sigma(), split_instrument_offset)
agree = abs(math.log(attn_a / attn_b)) <= threshold
ceilings = {budget: ms / (budget * attn_a + identity_ms) for budget in budgets} if agree else None
split[config_id] = dict(ms=ms, rest_ms=identity_ms, attn_a_ms=attn_a, attn_b_ms=attn_b, agree=agree, ceilings=ceilings)
agree_cell = "-" if agree is None else ("yes" if agree else "[yellow]no[/yellow]")
ceiling_cells = [f"{ceilings[budget]:.2f}x" if ceilings else "-" for budget in budgets]
table.add_row(entry["label"], f"{ms:8.3f} ms", f"{identity_ms:8.3f} ms", f"{attn_a:8.3f} ms", f"{attn_b:8.3f} ms" if attn_b else "-", agree_cell, *ceiling_cells)
if split:
subtitle = "[dim]rest = identity-attention variant; A = block minus rest, B = standalone attention table"
subtitle += f" ({attention_preset})" if attention_preset else " (no matching attention preset this run)"
subtitle += "; ceilings are per-block upper bounds at the given kv budget, generation adds te/vae/projections[/dim]"
emit(Panel(table, title=f"compute split: {family}", subtitle=subtitle, box=ROUNDED_BOX, expand=False))
disagreements = [config_id for config_id, row in split.items() if row["agree"] is False]
if disagreements:
emit(f"[yellow]split instruments disagree on {', '.join(disagreements)}; ceilings withheld there, treat the split with suspicion[/yellow]")
return split
def measured(results, config_id):
entry = results.get(config_id) or {}
ms = entry.get("ms")
@@ -2661,7 +2933,7 @@ def measured(results, config_id):
def best_config(results):
# lowest error among rows within 5% of the fastest sdnq time
candidates = [(config_id, entry["ms"], entry["err"]) for config_id, entry in results.items() if entry.get("ms") is not None and config_id not in external_config_ids and config_id not in unsafe_config_ids]
candidates = [(config_id, entry["ms"], entry["err"]) for config_id, entry in results.items() if entry.get("ms") is not None and config_id not in external_config_ids and config_id not in sparse_config_ids and config_id not in unsafe_config_ids]
if not candidates:
return None
fastest = min(ms for _config_id, ms, _err in candidates)
@@ -2695,7 +2967,7 @@ def select_attention_config(results):
pool, capped = [], []
for config_id, label, kwargs in bench_configs:
settings = config_settings(kwargs)
if settings is None:
if settings is None or config_id in sparse_config_ids: # a sparse row shares a settings tuple with its dense row but is a stage over it, not a setting
continue
if settings["accum"] and settings["pv"] == "disabled":
continue # the unsafe accumulation combo is never a candidate; the accum row cites it directly
@@ -3503,11 +3775,24 @@ def main():
sys.exit(1)
known_blocks = [config_id for config_id, _w, _mm, _a in block_configs]
selected_blocks = None if args.block_configs.strip().lower() == "all" else [s.strip() for s in args.block_configs.split(",") if s.strip()]
selected_geometries = None if args.block_geometries.strip().lower() == "all" else [s.strip() for s in args.block_geometries.split(",") if s.strip()]
if selected_geometries is not None:
unknown_geometries = [s for s in selected_geometries if s not in block_geometries]
if unknown_geometries:
console.print(f"[red]unknown block geometry(ies): {', '.join(unknown_geometries)}; available: {', '.join(block_geometries)}[/red]")
sys.exit(1)
if selected_blocks is not None:
unknown_blocks = [s for s in selected_blocks if s not in known_blocks]
if unknown_blocks:
console.print(f"[red]unknown block config(s): {', '.join(unknown_blocks)}; available: {', '.join(known_blocks)}[/red]")
sys.exit(1)
known_attention = [config_id for config_id, _label, _kwargs in bench_configs]
selected_attention = None if args.configs.strip().lower() == "all" else [s.strip() for s in args.configs.split(",") if s.strip()]
if selected_attention is not None:
unknown_attention = [s for s in selected_attention if s not in known_attention]
if unknown_attention:
console.print(f"[red]unknown attention config(s): {', '.join(unknown_attention)}; available: {', '.join(known_attention)}[/red]")
sys.exit(1)
known_variants = [variant_id for variant_id, _cfg in dequant_variant_configs]
variants_arg = args.dequant_variants.strip().lower()
if variants_arg == "all":
@@ -3550,7 +3835,9 @@ def main():
bench_dtype = devices.dtype
else:
bench_dtype = {"bf16": torch.bfloat16, "fp16": torch.float16}[args.dtype]
selected = list(full_run) if args.shapes.strip().lower() == "all" else [s.strip() for s in args.shapes.split(",") if s.strip()]
shapes_arg = args.shapes.strip().lower()
shape_run_aliases = {"all": full_run, "sparse": sparse_run, "gate": gate_run}
selected = list(shape_run_aliases[shapes_arg]) if shapes_arg in shape_run_aliases else [s.strip() for s in args.shapes.split(",") if s.strip()]
unknown = [s for s in selected if s not in shape_presets]
if unknown:
console.print(f"[red]unknown shape preset(s): {', '.join(unknown)}; available: {', '.join(shape_presets)}[/red]")
@@ -3615,7 +3902,9 @@ def main():
if free_vram_gb() < 3.0:
emit(f"[yellow]skipping block benchmarks: needs about 3 gb free vram, {free_vram_gb():.1f} gb available[/yellow]")
else:
bench_block_section(args.iters, args.warmup, config_timeout=args.config_timeout, selected=selected_blocks)
bench_block_section(args.iters, args.warmup, config_timeout=args.timeout_flag, selected=selected_blocks, geometries=selected_geometries)
if "attention" not in sections:
emit_block_splits()
if "attention" in sections:
# bench the prep mode the advice points to: compiled, static workaround, or eager
@@ -3637,7 +3926,8 @@ def main():
if free_vram_gb() < needed:
emit(f"[yellow]skipping {preset}: needs about {needed:.0f} gb free vram, {free_vram_gb():.1f} gb available[/yellow]")
continue
all_results[preset] = bench_shape(preset, args.iters, args.warmup, position=(index, len(selected)), config_timeout=args.config_timeout, fp8_result=fp8_result)
all_results[preset] = bench_shape(preset, args.iters, args.warmup, position=(index, len(selected)), config_timeout=args.timeout_flag, fp8_result=fp8_result, selected=selected_attention)
emit_block_splits()
build_recommendations(all_results, fp8_result, prep_status, block_results=(report.get("block") or {}).get("results"), block_variants=report.get("blocks"))
if drift_samples:
constraints.txt
-1
View File
@@ -1,3 +1,2 @@
fastapi==0.124.4
numpy==2.1.2
Pillow==12.2.0
data/reference-base.json
+776 -760
View File
File diff suppressed because it is too large Load Diff
data/reference-cloud.json
+15 -11
View File
@@ -1,22 +1,26 @@
{
"Nano Banana": {
"path": "gemini-2.5-flash-image",
"desc": "Our best engine for high-velocity visual creation, offering state-of-the-art speed and efficiency. Gemini 2.5 Flash Image, also known as Nano Banana, is best for high-volume generation, conversational image editing, and low-latency creative workflows that require native multimodal understanding. (Knowledge cutoff June 2025)",
"preview": "gemini-2.5-flash-image.jpg"
"Nano Banana lite": {
"path": "gemini-3.1-flash-lite-image",
"desc": "Nano Banana Lite is designed as the efficiency specialist of the image generation family, offering ultra-low latency and cost-effective image generation and editing. By targeting a sub-2 second latency and significantly reduced TPU compute costs, this model enables high-volume interactive developer use cases and real-time consumer applications. (Knowledge cutoff January 2025)",
"preview": "gemini-3.1-flash-lite-image.jpg",
"date": "2026 July"
},
"Nano Banana 2": {
"path": "gemini-3.1-flash-image",
"desc": "Nano Banana 2 provides high-quality image generation and conversational editing at a mainstream price point and low latency. It serves as the high-efficiency counterpart to Gemini 3 Pro Image, optimized for speed and high-volume developer use cases.(Knowledge cutoff January 2025)",
"preview": "gemini-3.1-flash-image.jpg"
},
"Nano Banana lite": {
"path": "gemini-3.1-flash-lite-image",
"desc": "Nano Banana Lite is designed as the efficiency specialist of the image generation family, offering ultra-low latency and cost-effective image generation and editing. By targeting a sub-2 second latency and significantly reduced TPU compute costs, this model enables high-volume interactive developer use cases and real-time consumer applications. (Knowledge cutoff January 2025)",
"preview": "gemini-3.1-flash-lite-image.jpg"
"preview": "gemini-3.1-flash-image.jpg",
"date": "2026 February"
},
"Nano Banana Pro": {
"path": "gemini-3-pro-image",
"desc": "Nano Banana Pro is a sophisticated reasoning-driven engine for professional-grade image editing and generation, offering studio-quality precision and advanced creative control. Nano Banana Pro is best for complex graphic design, high-fidelity product mockups, and factual data visualizations that require accurate text rendering and real-world grounding via Google Search. (Knowledge cutoff January 2025)",
"preview": "gemini-3-pro-image.jpg"
"preview": "gemini-3-pro-image.jpg",
"date": "2025 November"
},
"Nano Banana": {
"path": "gemini-2.5-flash-image",
"desc": "Our best engine for high-velocity visual creation, offering state-of-the-art speed and efficiency. Gemini 2.5 Flash Image, also known as Nano Banana, is best for high-volume generation, conversational image editing, and low-latency creative workflows that require native multimodal understanding. (Knowledge cutoff June 2025)",
"preview": "gemini-2.5-flash-image.jpg",
"date": "2025 August"
}
}
data/reference-community.json
+99 -99
View File
@@ -1,17 +1,12 @@
{
"Tempest-by-Vlad XL": {
"path": "tempestByVlad_baseV01.safetensors@https://civitai.com/api/download/models/1301775",
"preview": "tempestByVlad_baseV01.jpg",
"desc": "Flexible SDXL model with custom encoder and finetuned for larger landscape resolutions with high details and high contrast.",
"size": 6.94,
"date": "2025 January"
},
"Tempest-by-Vlad XL Hyper": {
"path": "tempestByVlad_hyperV01.safetensors@https://civitai.com/api/download/models/1343512",
"preview": "tempestByVlad_hyperV01.jpg",
"desc": "Custom distilled variant with goal to get as-normal-as-possible model that works with low steps and guidance-free",
"size": 6.94,
"date": "2025 January"
"Juggernaut SD Reborn": {
"original": true,
"path": "juggernaut_reborn.safetensors@https://civitai.com/api/download/models/274039",
"preview": "juggernaut_reborn.jpg",
"desc": "Showcase finetuned model based on Stable diffusion 1.5",
"date": "2023 December",
"size": 2.28,
"extras": "width: 512, height: 512, sampler: DEIS, steps: 20, cfg_scale: 6.0"
},
"Juggernaut XL XI": {
"path": "juggernautXL_juggXIByRundiffusion.safetensors@https://civitai.com/api/download/models/782002",
@@ -29,28 +24,26 @@
"size": 6.94,
"extras": "sampler: DPM SDE, steps: 6, cfg_scale: 2.0"
},
"Juggernaut SD Reborn": {
"original": true,
"path": "juggernaut_reborn.safetensors@https://civitai.com/api/download/models/274039",
"preview": "juggernaut_reborn.jpg",
"desc": "Showcase finetuned model based on Stable diffusion 1.5",
"date": "2023 December",
"size": 2.28,
"extras": "width: 512, height: 512, sampler: DEIS, steps: 20, cfg_scale: 6.0"
},
"WAI Illustrious XL v15": {
"path": "waiIllustriousSDXL_v150.safetensors@https://civitai.com/api/download/models/2167369",
"preview": "waiIllustriousSDXL_v150.jpg",
"NoobAI XL 1.1 Epsilon": {
"path": "noobaiXLNAIXL_epsilonPred11Version.safetensors@https://huggingface.co/Laxhar/noobai-XL-1.1/resolve/main/NoobAI-XL-v1.1.safetensors",
"preview": "noobaiXLNAIXL_epsilonPred11Version.jpg",
"desc": "",
"size": 6.94,
"date": "2025 August"
"date": "2024 November"
},
"Pony Realism XL v2.3": {
"path": "ponyRealism_V23.safetensors@https://civitai.com/api/download/models/1763661",
"preview": "ponyRealism_V23.jpg",
"desc": "",
"size": 6.94,
"date": "2025 May"
"ShuttleAI Shuttle 3.0 Diffusion": {
"path": "shuttleai/shuttle-3-diffusion",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-3-diffusion.jpg",
"date": "2024 November",
"size": 31.41
},
"ShuttleAI Shuttle 3.1 Aesthetic": {
"path": "shuttleai/shuttle-3.1-aesthetic",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-3.1-aesthetic.jpg",
"date": "2024 November",
"size": 31.41
},
"NoobAI XL 1.0 V-Pred": {
"path": "noobaiXLNAIXL_vPred10Version.safetensors@https://huggingface.co/Laxhar/noobai-XL-Vpred-1.0/resolve/main/NoobAI-XL-Vpred-v1.0.safetensors",
@@ -59,12 +52,33 @@
"size": 6.94,
"date": "2024 December"
},
"NoobAI XL 1.1 Epsilon": {
"path": "noobaiXLNAIXL_epsilonPred11Version.safetensors@https://huggingface.co/Laxhar/noobai-XL-1.1/resolve/main/NoobAI-XL-v1.1.safetensors",
"preview": "noobaiXLNAIXL_epsilonPred11Version.jpg",
"Tempest-by-Vlad XL": {
"path": "tempestByVlad_baseV01.safetensors@https://civitai.com/api/download/models/1301775",
"preview": "tempestByVlad_baseV01.jpg",
"desc": "Flexible SDXL model with custom encoder and finetuned for larger landscape resolutions with high details and high contrast.",
"size": 6.94,
"date": "2025 January"
},
"Tempest-by-Vlad XL Hyper": {
"path": "tempestByVlad_hyperV01.safetensors@https://civitai.com/api/download/models/1343512",
"preview": "tempestByVlad_hyperV01.jpg",
"desc": "Custom distilled variant with goal to get as-normal-as-possible model that works with low steps and guidance-free",
"size": 6.94,
"date": "2025 January"
},
"ShuttleAI Shuttle Jaguar": {
"path": "shuttleai/shuttle-jaguar",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-jaguar.jpg",
"date": "2025 January",
"size": 31.41
},
"Pony Realism XL v2.3": {
"path": "ponyRealism_V23.safetensors@https://civitai.com/api/download/models/1763661",
"preview": "ponyRealism_V23.jpg",
"desc": "",
"size": 6.94,
"date": "2024 November"
"date": "2025 May"
},
"WAI-Ani-Pony XL v14": {
"path": "waiANIPONYXL_v140.safetensors@https://civitai.com/api/download/models/1767402",
@@ -73,6 +87,55 @@
"size": 6.94,
"date": "2025 May"
},
"WAI Illustrious XL v15": {
"path": "waiIllustriousSDXL_v150.safetensors@https://civitai.com/api/download/models/2167369",
"preview": "waiIllustriousSDXL_v150.jpg",
"desc": "",
"size": 6.94,
"date": "2025 August"
},
"Tiwaz CenKreChro": {
"path": "Tiwaz/CenKreChro",
"preview": "Tiwaz--CenKreChro.jpg",
"desc": "Based Centerfold Flux 5, trying to merge in Chroma and Krea.",
"date": "2025 September",
"size": 31.42
},
"purplesmartai Pony 7": {
"path": "purplesmartai/pony-v7-base",
"preview": "purplesmartai--pony-v7-base.jpg",
"desc": "Pony V7 is a versatile character generation model based on AuraFlow architecture. It supports a wide range of styles and species types (humanoid, anthro, feral, and more) and handles character interactions through natural language prompts.",
"date": "2025 October",
"size": 33.32
},
"Skywork UniPic3": {
"path": "Skywork/Unipic3",
"preview": "Skywork--Unipic3.jpg",
"desc": "UniPic3 is an image editing and multi-image composition model based. It is a fine-tune of Qwen-Image-Edit.",
"date": "2026 January",
"size": 53.74
},
"Skywork Unipic3-DMD": {
"path": "Skywork/Unipic3-DMD",
"preview": "Skywork--Unipic3-DMD.jpg",
"desc": "UniPic3-DMD-Model is a few-step image editing and multi-image composition model trained using Distribution Matching Distillation (DMD) and is a fine-tune of Qwen-Image-Edit.",
"date": "2026 January",
"size": 53.74
},
"FireRed Image Edit 1.0": {
"path": "FireRedTeam/FireRed-Image-Edit-1.0",
"preview": "FireRedTeam--FireRed-Image-Edit-1.0.jpg",
"desc": "FireRed-Image-Edit is a general-purpose image editing model that delivers high-fidelity and consistent editing across a wide range of scenarios. FireRed is a fine-tune of Qwen-Image-Edit.",
"date": "2026 February",
"size": 53.74
},
"FireRed Image Edit 1.1": {
"path": "FireRedTeam/FireRed-Image-Edit-1.1",
"preview": "FireRedTeam--FireRed-Image-Edit-1.1.jpg",
"desc": "FireRed-Image-Edit is a general-purpose image editing model that delivers high-fidelity and consistent editing across a wide range of scenarios. FireRed is a fine-tune of Qwen-Image-Edit.",
"date": "2026 March",
"size": 53.74
},
"Z-Image-Turbo MoodyRealMix": {
"path": "resonantsky/MoodyRealMix-SDNQ-int8-svd-r32",
"preview": "resonantsky--MoodyRealMix-SDNQ-int8-svd-r32.jpg",
@@ -100,69 +163,6 @@
"tags": "community, Z-image",
"date": "2026 May"
},
"Tiwaz CenKreChro": {
"path": "Tiwaz/CenKreChro",
"preview": "Tiwaz--CenKreChro.jpg",
"desc": "Based Centerfold Flux 5, trying to merge in Chroma and Krea.",
"date": "2025 September",
"size": 31.42
},
"purplesmartai Pony 7": {
"path": "purplesmartai/pony-v7-base",
"preview": "purplesmartai--pony-v7-base.jpg",
"desc": "Pony V7 is a versatile character generation model based on AuraFlow architecture. It supports a wide range of styles and species types (humanoid, anthro, feral, and more) and handles character interactions through natural language prompts.",
"date": "2025 October",
"size": 33.32
},
"ShuttleAI Shuttle 3.0 Diffusion": {
"path": "shuttleai/shuttle-3-diffusion",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-3-diffusion.jpg",
"date": "2024 November",
"size": 31.41
},
"ShuttleAI Shuttle 3.1 Aesthetic": {
"path": "shuttleai/shuttle-3.1-aesthetic",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-3.1-aesthetic.jpg",
"date": "2024 November",
"size": 31.41
},
"ShuttleAI Shuttle Jaguar": {
"path": "shuttleai/shuttle-jaguar",
"desc": "Shuttle uses Flux.1 Schnell as its base. It can produce images similar to Flux Dev or Pro in just 4 steps, and it is licensed under Apache 2. The model was partially de-distilled during training. When used beyond 10 steps, it enters refiner mode enhancing image details without altering the composition",
"preview": "shuttleai--shuttle-jaguar.jpg",
"date": "2025 January",
"size": 31.41
},
"FireRed Image Edit 1.0": {
"path": "FireRedTeam/FireRed-Image-Edit-1.0",
"preview": "FireRedTeam--FireRed-Image-Edit-1.0.jpg",
"desc": "FireRed-Image-Edit is a general-purpose image editing model that delivers high-fidelity and consistent editing across a wide range of scenarios. FireRed is a fine-tune of Qwen-Image-Edit.",
"date": "2026 February",
"size": 53.74
},
"FireRed Image Edit 1.1": {
"path": "FireRedTeam/FireRed-Image-Edit-1.1",
"preview": "FireRedTeam--FireRed-Image-Edit-1.1.jpg",
"desc": "FireRed-Image-Edit is a general-purpose image editing model that delivers high-fidelity and consistent editing across a wide range of scenarios. FireRed is a fine-tune of Qwen-Image-Edit.",
"date": "2026 March",
"size": 53.74
},
"Skywork UniPic3": {
"path": "Skywork/Unipic3",
"preview": "Skywork--Unipic3.jpg",
"desc": "UniPic3 is an image editing and multi-image composition model based. It is a fine-tune of Qwen-Image-Edit.",
"date": "2026 January",
"size": 53.74
},
"Skywork Unipic3-DMD": {
"path": "Skywork/Unipic3-DMD",
"preview": "Skywork--Unipic3-DMD.jpg",
"desc": "UniPic3-DMD-Model is a few-step image editing and multi-image composition model trained using Distribution Matching Distillation (DMD) and is a fine-tune of Qwen-Image-Edit.",
"date": "2026 January",
"size": 53.74
},
"Anima 1.0 Base Merge sdnq-hadamard-uint4": {
"path": "vladmandic/Anima-1.0-Base-Merge-sdnq-hadamard-uint4",
"preview": "vladmandic--Anima-1.0-Base.jpg",
data/reference-distilled.json
+125 -116
View File
@@ -1,17 +1,20 @@
{
"Boogu Image 0.1 Turbo": {
"path": "Boogu/Boogu-Image-0.1-Turbo",
"preview": "Boogu--Boogu-Image-0.1-Turbo.jpg",
"desc": "Boogu Image 0.1 Turbo is the distilled fast inference variant of Boogu Image with the same Qwen3-VL instruction encoder and Boogu transformer architecture.",
"size": 35.81,
"date": "2026 June"
"Segmind Tiny": {
"path": "segmind/tiny-sd",
"preview": "segmind--tiny-sd.jpg",
"desc": "Segmind's Tiny-SD offers a compact, efficient, and distilled version of Realistic Vision 4.0 and is up to 80% faster than SD1.5",
"extras": "width: 512, height: 512, sampler: Default, cfg_scale: 9.0",
"size": 0.99,
"date": "2023 July"
},
"Boogu Image 0.1 Edit Turbo": {
"path": "Boogu/Boogu-Image-0.1-Edit-Turbo",
"preview": "Boogu--Boogu-Image-0.1-Edit-Turbo.jpg",
"desc": "Boogu Image 0.1 Edit Turbo is the distilled editing variant of Boogu Image with motion-aware instruction encoding and fast flow-match inference.",
"size": 35.81,
"date": "2026 June"
"Segmind SSD-1B": {
"path": "huggingface/segmind/SSD-1B",
"preview": "segmind--SSD-1B.jpg",
"desc": "The Segmind Stable Diffusion Model (SSD-1B) offers a compact, efficient, and distilled version of the SDXL model. At 50% smaller and 60% faster than Stable Diffusion XL (SDXL), it provides quick and seamless performance without sacrificing image quality.",
"variant": "fp16",
"extras": "sampler: Default, cfg_scale: 9.0",
"size": 12.48,
"date": "2023 October"
},
"StabilityAI StableDiffusion XL Turbo": {
"path": "stabilityai/sdxl-turbo",
@@ -22,13 +25,29 @@
"size": 19.38,
"date": "2023 November"
},
"Krea 2 Turbo": {
"path": "CalamitousFelicitousness/Krea-2-Turbo-Diffusers",
"preview": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
"desc": "Krea 2 (K2) Turbo is the 8-step distilled inference model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. Runs without classifier-free guidance; LoRAs trained on Krea 2 Base apply directly.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 33.5,
"date": "2026 June"
"SDXL Flash Mini": {
"path": "SDXL-Flash_Mini.safetensors@https://huggingface.co/sd-community/sdxl-flash-mini/resolve/main/SDXL-Flash_Mini.safetensors?download=true",
"preview": "SDXL-Flash_Mini.jpg",
"desc": "Introducing the new fast model SDXL Flash (Mini), we learned that all fast XL models work fast, but the quality decreases, and we also made a fast model, but it is not as fast as LCM, Turbo, Lightning and Hyper, but the quality is higher.",
"extras": "sampler: DEIS, steps: 40, cfg_scale: 6.0",
"experimental": true,
"date": "2024 May"
},
"Tencent HunyuanDiT 1.1 Distilled": {
"path": "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled",
"desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
"preview": "Tencent-Hunyuan--HunyuanDiT-v1.1-Diffusers-Distilled.jpg",
"extras": "sampler: Default, cfg_scale: 2.0",
"size": 13.49,
"date": "2024 June"
},
"Tencent HunyuanDiT 1.2 Distilled": {
"path": "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers-Distilled",
"desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
"preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers-Distilled.jpg",
"extras": "sampler: Default, cfg_scale: 2.0",
"size": 13.43,
"date": "2024 July"
},
"StabilityAI Stable Diffusion 3.5 Turbo": {
"path": "stabilityai/stable-diffusion-3.5-large-turbo",
@@ -39,28 +58,12 @@
"size": 36.12,
"date": "2024 October"
},
"Microsoft Lens Turbo": {
"path": "Jinstudio/Lens-Turbo",
"preview": "microsoft--Lens-Turbo.jpg",
"desc": "Microsoft Lens-Turbo is the distilled Lens variant optimized for faster text-to-image generation with fewer steps.",
"size": 28.43,
"date": "2026 May"
},
"Tencent FLUX.1 Dev SRPO": {
"path": "vladmandic/flux.1-dev-SRPO",
"preview": "vladmandic--flux.1-dev-SRPO.jpg",
"desc": "FLUX.1 Dev SRPO is Tencent trained with specific technique: Directly Aligning the Full Diffusion Trajectory with Fine-Grained Human Preference",
"extras": "sampler: Default, cfg_scale: 4.5",
"size": 31.42,
"date": "2025 September"
},
"HiDream-O1 Image Dev": {
"path": "HiDream-ai/HiDream-O1-Image-Dev",
"preview": "HiDream-ai--HiDream-O1-Image-Dev.jpg",
"desc": "HiDream-O1-Image-Dev is the distilled 8B HiDream-O1 variant tuned for 28-step fast generation using flash flow scheduling.",
"extras": "sampler: Flash, steps: 28, cfg_scale: 0.0",
"size": 35.2,
"date": "2026 May"
"NVLabs Sana 1.5 1.6B 1k Sprint": {
"path": "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
"desc": "SANA-Sprint is an ultra-efficient diffusion model for text-to-image (T2I) generation, reducing inference steps from 20 to 1-4 while achieving state-of-the-art performance.",
"preview": "Efficient-Large-Model--Sana15_Sprint_1600M_1024px_diffusers.jpg",
"size": 9.03,
"date": "2025 March"
},
"Qwen-Image-Lightning": {
"path": "vladmandic/Qwen-Lightning",
@@ -78,13 +81,20 @@
"size": 56.1,
"date": "2025 August"
},
"Baidu ERNIE-Image-Turbo": {
"path": "baidu/ERNIE-Image-Turbo",
"preview": "baidu--ERNIE-Image-Turbo.jpg",
"desc": "ERNIE-Image-Turbo is a distilled ERNIE-Image variant optimized for fast generation with fewer denoising steps.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 22.29,
"date": "2026 April"
"lodestones Chroma1 Flash": {
"path": "lodestones/Chroma1-Flash",
"preview": "lodestones--Chroma1-Flash.jpg",
"desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. A fine-tuned version of the Chroma1-Base made to find the best way to make these flow matching models faster.",
"size": 25.6,
"date": "2025 August"
},
"Tencent FLUX.1 Dev SRPO": {
"path": "vladmandic/flux.1-dev-SRPO",
"preview": "vladmandic--flux.1-dev-SRPO.jpg",
"desc": "FLUX.1 Dev SRPO is Tencent trained with specific technique: Directly Aligning the Full Diffusion Trajectory with Fine-Grained Human Preference",
"extras": "sampler: Default, cfg_scale: 4.5",
"size": 31.42,
"date": "2025 September"
},
"Qwen-Image-Lightning-Edit": {
"path": "vladmandic/Qwen-Lightning-Edit",
@@ -110,6 +120,13 @@
"date": "2025 September",
"size": 41.08
},
"Tencent HunyuanImage 2.1 Distilled": {
"path": "hunyuanvideo-community/HunyuanImage-2.1-Distilled-Diffusers",
"desc": "HunyuanImage-2.1, a highly efficient text-to-image model that is capable of generating 2K (2048 × 2048) resolution images.",
"preview": "hunyuanvideo-community--HunyuanImage-2.1-Distilled-Diffusers.jpg",
"size": 49.53,
"date": "2025 September"
},
"Qwen-Image-Edit-2509 Pruning-13B": {
"path": "OPPOer/Qwen-Image-Edit-2509-Pruning",
"subfolder": "Qwen-Image-Edit-2509-13B-4steps",
@@ -118,51 +135,6 @@
"date": "2025 October",
"size": 42.34
},
"lodestones Chroma1 Flash": {
"path": "lodestones/Chroma1-Flash",
"preview": "lodestones--Chroma1-Flash.jpg",
"desc": "Chroma is a 8.9B parameter model based on FLUX.1-schnell. It’s fully Apache 2.0 licensed, ensuring that anyone can use, modify, and build on top of it—no corporate gatekeeping. A fine-tuned version of the Chroma1-Base made to find the best way to make these flow matching models faster.",
"size": 25.6,
"date": "2025 August"
},
"SDXL Flash Mini": {
"path": "SDXL-Flash_Mini.safetensors@https://huggingface.co/sd-community/sdxl-flash-mini/resolve/main/SDXL-Flash_Mini.safetensors?download=true",
"preview": "SDXL-Flash_Mini.jpg",
"desc": "Introducing the new fast model SDXL Flash (Mini), we learned that all fast XL models work fast, but the quality decreases, and we also made a fast model, but it is not as fast as LCM, Turbo, Lightning and Hyper, but the quality is higher.",
"extras": "sampler: DEIS, steps: 40, cfg_scale: 6.0",
"experimental": true
},
"NVLabs Sana 1.5 1.6B 1k Sprint": {
"path": "Efficient-Large-Model/Sana_Sprint_1.6B_1024px_diffusers",
"desc": "SANA-Sprint is an ultra-efficient diffusion model for text-to-image (T2I) generation, reducing inference steps from 20 to 1-4 while achieving state-of-the-art performance.",
"preview": "Efficient-Large-Model--Sana15_Sprint_1600M_1024px_diffusers.jpg",
"size": 9.03,
"date": "2025 March"
},
"Segmind SSD-1B": {
"path": "huggingface/segmind/SSD-1B",
"preview": "segmind--SSD-1B.jpg",
"desc": "The Segmind Stable Diffusion Model (SSD-1B) offers a compact, efficient, and distilled version of the SDXL model. At 50% smaller and 60% faster than Stable Diffusion XL (SDXL), it provides quick and seamless performance without sacrificing image quality.",
"variant": "fp16",
"extras": "sampler: Default, cfg_scale: 9.0",
"size": 12.48,
"date": "2023 October"
},
"Segmind Tiny": {
"path": "segmind/tiny-sd",
"preview": "segmind--tiny-sd.jpg",
"desc": "Segmind's Tiny-SD offers a compact, efficient, and distilled version of Realistic Vision 4.0 and is up to 80% faster than SD1.5",
"extras": "width: 512, height: 512, sampler: Default, cfg_scale: 9.0",
"size": 0.99,
"date": "2023 July"
},
"Tencent HunyuanImage 2.1 Distilled": {
"path": "hunyuanvideo-community/HunyuanImage-2.1-Distilled-Diffusers",
"desc": "HunyuanImage-2.1, a highly efficient text-to-image model that is capable of generating 2K (2048 × 2048) resolution images.",
"preview": "hunyuanvideo-community--HunyuanImage-2.1-Distilled-Diffusers.jpg",
"size": 49.53,
"date": "2025 September"
},
"Bria Fibo-Lite": {
"path": "briaai/Fibo-lite",
"preview": "briaai--Fibo-lite.jpg",
@@ -171,22 +143,6 @@
"size": 22.47,
"date": "2025 November"
},
"Tencent HunyuanDiT 1.2 Distilled": {
"path": "Tencent-Hunyuan/HunyuanDiT-v1.2-Diffusers-Distilled",
"desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
"preview": "Tencent-Hunyuan--HunyuanDiT-v1.2-Diffusers-Distilled.jpg",
"extras": "sampler: Default, cfg_scale: 2.0",
"size": 13.43,
"date": "2024 July"
},
"Tencent HunyuanDiT 1.1 Distilled": {
"path": "Tencent-Hunyuan/HunyuanDiT-v1.1-Diffusers-Distilled",
"desc": "Hunyuan-DiT : A Powerful Multi-Resolution Diffusion Transformer with Fine-Grained Chinese Understanding.",
"preview": "Tencent-Hunyuan--HunyuanDiT-v1.1-Diffusers-Distilled.jpg",
"extras": "sampler: Default, cfg_scale: 2.0",
"size": 13.49,
"date": "2024 June"
},
"Black Forest Labs FLUX.2 Klein 4B": {
"path": "black-forest-labs/FLUX.2-klein-4B",
"preview": "black-forest-labs--FLUX.2-klein-4B.jpg",
@@ -203,6 +159,13 @@
"size": 32.32,
"date": "2026 January"
},
"Meituan LongCat Image-Edit Turbo": {
"path": "meituan-longcat/LongCat-Image-Edit-Turbo",
"preview": "meituan-longcat--LongCat-Image-Edit.jpg",
"desc": "LongCat-Image-Edit-Turbo, the distilled version of LongCat-Image-Edit. It achieves high-quality image editing with only 8 NFEs (Number of Function Evaluations) , offering extremely low inference latency.",
"size": 27.28,
"date": "2026 February"
},
"Black Forest Labs FLUX.2 Klein 9B KV": {
"path": "black-forest-labs/FLUX.2-klein-9b-kv",
"preview": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
@@ -211,6 +174,51 @@
"size": 32.32,
"date": "2026 March"
},
"Baidu ERNIE-Image-Turbo": {
"path": "baidu/ERNIE-Image-Turbo",
"preview": "baidu--ERNIE-Image-Turbo.jpg",
"desc": "ERNIE-Image-Turbo is a distilled ERNIE-Image variant optimized for fast generation with fewer denoising steps.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 22.29,
"date": "2026 April"
},
"Microsoft Lens Turbo": {
"path": "Jinstudio/Lens-Turbo",
"preview": "microsoft--Lens-Turbo.jpg",
"desc": "Microsoft Lens-Turbo is the distilled Lens variant optimized for faster text-to-image generation with fewer steps.",
"size": 28.43,
"date": "2026 May"
},
"HiDream-O1 Image Dev": {
"path": "HiDream-ai/HiDream-O1-Image-Dev",
"preview": "HiDream-ai--HiDream-O1-Image-Dev.jpg",
"desc": "HiDream-O1-Image-Dev is the distilled 8B HiDream-O1 variant tuned for 28-step fast generation using flash flow scheduling.",
"extras": "sampler: Flash, steps: 28, cfg_scale: 0.0",
"size": 35.2,
"date": "2026 May"
},
"Krea 2 Turbo": {
"path": "CalamitousFelicitousness/Krea-2-Turbo-Diffusers",
"preview": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
"desc": "Krea 2 (K2) Turbo is the 8-step distilled inference model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. Runs without classifier-free guidance; LoRAs trained on Krea 2 Base apply directly.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 33.5,
"date": "2026 June"
},
"Boogu Image 0.1 Turbo": {
"path": "Boogu/Boogu-Image-0.1-Turbo",
"preview": "Boogu--Boogu-Image-0.1-Turbo.jpg",
"desc": "Boogu Image 0.1 Turbo is the distilled fast inference variant of Boogu Image with the same Qwen3-VL instruction encoder and Boogu transformer architecture.",
"size": 35.81,
"date": "2026 June"
},
"Boogu Image 0.1 Edit Turbo": {
"path": "Boogu/Boogu-Image-0.1-Edit-Turbo",
"preview": "Boogu--Boogu-Image-0.1-Edit-Turbo.jpg",
"desc": "Boogu Image 0.1 Edit Turbo is the distilled editing variant of Boogu Image with motion-aware instruction encoding and fast flow-match inference.",
"size": 35.81,
"date": "2026 June"
},
"Anima 1.0 Turbo": {
"path": "CalamitousFelicitousness/Anima-1.0-Turbo-Diffusers",
"preview": "CalamitousFelicitousness--Anima-1.0-Turbo-Diffusers.jpg",
@@ -219,13 +227,6 @@
"date": "2026 July",
"size": 4.99
},
"Meituan LongCat Image-Edit Turbo": {
"path": "meituan-longcat/LongCat-Image-Edit-Turbo",
"preview": "meituan-longcat--LongCat-Image-Edit.jpg",
"desc": "LongCat-Image-Edit-Turbo, the distilled version of LongCat-Image-Edit. It achieves high-quality image editing with only 8 NFEs (Number of Function Evaluations) , offering extremely low inference latency.",
"size": 27.28,
"date": "2026 February"
},
"Microsoft Mage-Flow Turbo": {
"path": "vladmandic/Mage-Flow-4B-Turbo",
"preview": "vladmandic--Mage-Flow-Turbo-4B.jpg",
@@ -257,5 +258,13 @@
"extras": "sampler: Default",
"size": 17.69,
"date": "2026 July"
},
"inclusionAI LLaDA-Image Turbo": {
"path": "inclusionAI/LLaDA-Image-Turbo",
"preview": "inclusionAI--LLaDA-Image-Turbo.jpg",
"desc": "LLaDA-Image-Turbo is the distilled fast-generation and editing variant of LLaDA-Image.",
"extras": "steps: 4, cfg_scale: 1.0",
"size": 37.15,
"date": "2026 September"
}
}
data/reference-nunchaku.json
+65 -58
View File
@@ -1,4 +1,27 @@
{
"SDXL Base Nunchaku SVDQuant": {
"path": "stabilityai/stable-diffusion-xl-base-1.0",
"subfolder": "nunchaku",
"preview": "stabilityai--stable-diffusion-xl-base-1.0.jpg",
"desc": "Nunchaku SVDQuant quantization of SDXL Base 1.0 UNet with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"size": 32.0,
"date": "2023 July"
},
"SDXL Turbo Nunchaku SVDQuant": {
"path": "stabilityai/sdxl-turbo",
"subfolder": "nunchaku",
"preview": "stabilityai--sdxl-turbo.jpg",
"desc": "Nunchaku SVDQuant quantization of SDXL Turbo UNet with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 19.38,
"date": "2023 November"
},
"FLUX.1-Dev Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-dev",
"subfolder": "nunchaku",
@@ -24,30 +47,6 @@
"size": 31.41,
"date": "2024 July"
},
"FLUX.1-Kontext Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-Kontext-dev",
"subfolder": "nunchaku",
"preview": "black-forest-labs--FLUX.1-Kontext-dev.jpg",
"desc": "Nunchaku SVDQuant quantization of FLUX.1-Kontext-dev transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model",
"TE"
],
"size": 31.42,
"date": "2025 May"
},
"FLUX.1-Krea Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-Krea-dev",
"subfolder": "nunchaku",
"preview": "black-forest-labs--FLUX.1-Krea-dev.jpg",
"desc": "Nunchaku SVDQuant quantization of FLUX.1-Krea-dev transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model",
"TE"
],
"size": 31.42,
"date": "2025 July"
},
"FLUX.1-Fill Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-Fill-dev",
"subfolder": "nunchaku",
@@ -74,6 +73,17 @@
"size": 40.68,
"date": "2024 November"
},
"Sana 1.6B 1k Nunchaku SVDQuant": {
"path": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
"subfolder": "nunchaku",
"preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"desc": "Nunchaku SVDQuant quantization of Sana 1.6B 1024px transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"size": 22.22,
"date": "2024 December"
},
"Shuttle Jaguar Nunchaku SVDQuant": {
"path": "shuttleai/shuttle-jaguar",
"subfolder": "nunchaku",
@@ -86,6 +96,30 @@
"size": 31.41,
"date": "2025 January"
},
"FLUX.1-Kontext Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-Kontext-dev",
"subfolder": "nunchaku",
"preview": "black-forest-labs--FLUX.1-Kontext-dev.jpg",
"desc": "Nunchaku SVDQuant quantization of FLUX.1-Kontext-dev transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model",
"TE"
],
"size": 31.42,
"date": "2025 May"
},
"FLUX.1-Krea Nunchaku SVDQuant": {
"path": "black-forest-labs/FLUX.1-Krea-dev",
"subfolder": "nunchaku",
"preview": "black-forest-labs--FLUX.1-Krea-dev.jpg",
"desc": "Nunchaku SVDQuant quantization of FLUX.1-Krea-dev transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model",
"TE"
],
"size": 31.42,
"date": "2025 July"
},
"Qwen-Image Nunchaku SVDQuant": {
"path": "Qwen/Qwen-Image",
"subfolder": "nunchaku",
@@ -167,17 +201,6 @@
"size": 53.74,
"date": "2025 September"
},
"Sana 1.6B 1k Nunchaku SVDQuant": {
"path": "Efficient-Large-Model/Sana_1600M_1024px_BF16_diffusers",
"subfolder": "nunchaku",
"preview": "Efficient-Large-Model--Sana_1600M_1024px_diffusers.jpg",
"desc": "Nunchaku SVDQuant quantization of Sana 1.6B 1024px transformer with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"size": 22.22,
"date": "2024 December"
},
"Z-Image-Turbo Nunchaku SVDQuant": {
"path": "Tongyi-MAI/Z-Image-Turbo",
"subfolder": "nunchaku",
@@ -190,29 +213,6 @@
"size": 30.58,
"date": "2025 November"
},
"SDXL Base Nunchaku SVDQuant": {
"path": "stabilityai/stable-diffusion-xl-base-1.0",
"subfolder": "nunchaku",
"preview": "stabilityai--stable-diffusion-xl-base-1.0.jpg",
"desc": "Nunchaku SVDQuant quantization of SDXL Base 1.0 UNet with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"size": 32.0,
"date": "2023 July"
},
"SDXL Turbo Nunchaku SVDQuant": {
"path": "stabilityai/sdxl-turbo",
"subfolder": "nunchaku",
"preview": "stabilityai--sdxl-turbo.jpg",
"desc": "Nunchaku SVDQuant quantization of SDXL Turbo UNet with INT4 and SVD rank 32",
"nunchaku": [
"Model"
],
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 19.38,
"date": "2023 November"
},
"Z-Image-Turbo Nunchaku-Lite": {
"path": "lite-infer/z-image-turbo-nunchaku-lite-int4_r32-bnb4-text-encoder",
"preview": "Tongyi-MAI--Z-Image-Turbo.jpg",
@@ -295,5 +295,12 @@
"desc": "Nunchaku-Lite quantization using precompiled Kernels",
"size": 10.92,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Nunchaku-Lite": {
"path": "rootonchair/MiniMax-H3-nunchaku-lite-int4",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"desc": "Nunchaku-Lite quantization using precompiled Kernels",
"size": 23.70,
"date": "2026 September"
}
}
data/reference-quantized.json
+135 -101
View File
@@ -27,37 +27,6 @@
"size": 12.6,
"date": "2025 October"
},
"FLUX.2 Dev sdnq-svd-uint4": {
"path": "Disty0/FLUX.2-dev-SDNQ-uint4-svd-r32",
"preview": "Disty0--FLUX.2-dev-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of black-forest-labs/FLUX.2-dev using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"size": 31.89,
"date": "2025 November"
},
"Black Forest Labs FLUX.2 Klein 4B sdnq-uint4-dynamic": {
"path": "Disty0/FLUX.2-klein-4B-SDNQ-4bit-dynamic",
"preview": "Disty0--FLUX.2-klein-4B-SDNQ-4bit-dynamic.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-4B using SDNQ.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 5.09,
"date": "2026 January"
},
"Black Forest Labs FLUX.2 Klein 9B sdnq-uint4-dynamic-svd": {
"path": "Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32",
"preview": "Disty0--FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-9B using SDNQ with SVD rank 32.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 11.73,
"date": "2026 January"
},
"Black Forest Labs FLUX.2 Klein 9B KV sdnq-uint4-dynamic-svd": {
"path": "vladmandic/Flux.2-Klein-9B-KV-sdnq-hadamard-uint4",
"preview": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-9B-KV using SDNQ with Hadamard.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 12.26,
"date": "2026 July"
},
"Chroma1-HD sdnq-svd-uint4": {
"path": "Disty0/Chroma1-HD-SDNQ-uint4-svd-r32",
"preview": "Disty0--Chroma1-HD-SDNQ-uint4-svd-r32.jpg",
@@ -79,48 +48,6 @@
"date": "2025 October",
"size": 23.53
},
"MiniMaxAI MiniMax-H3 sdnq-uint4": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_4bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 64.80,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 sdnq-uint4 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_4bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 64.80,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint4": {
"path": "OzzyGT/MiniMax_H3_sdnq_4bit_pruned",
"preview": "OzzyGT--MiniMax_H3_sdnq_4bit_pruned.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 23.70,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint4 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_4bit_pruned",
"preview": "OzzyGT--MiniMax_H3_sdnq_4bit_pruned.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 23.70,
"date": "2026 August"
},
"Z-Image-Turbo sdnq-svd-uint4": {
"path": "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
"preview": "Disty0--Z-Image-Turbo-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Tongyi-MAI/Z-Image-Turbo using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 9",
"size": 6.05,
"date": "2025 November"
},
"Qwen-Image sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-SDNQ-uint4-svd-r32",
"preview": "Qwen--Qwen-Image.jpg",
@@ -128,13 +55,6 @@
"date": "2025 October",
"size": 16.09
},
"Qwen-Image-2512 sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-2512-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-2512 using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"size": 16.09,
"date": "2026 January"
},
"Qwen-Image-Edit sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-Edit-SDNQ-uint4-svd-r32",
"preview": "Qwen--Qwen-Image-Edit.jpg",
@@ -149,20 +69,6 @@
"date": "2025 October",
"size": 16.09
},
"Qwen-Image-Edit-2511 sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-Edit-2511-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-Edit-2511-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-Edit-2511 using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"date": "2025 December",
"size": 16.09
},
"Qwen-Image-Layered sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-Layered-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-Layered-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-Layered using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"date": "2025 December",
"size": 16.09
},
"nVidia ChronoEdit sdnq-svd-uint4": {
"path": "Disty0/ChronoEdit-14B-SDNQ-uint4-svd-r32",
"preview": "Disty0--ChronoEdit-14B-SDNQ-uint4-svd-r32.jpg",
@@ -198,6 +104,58 @@
"size": 3.37,
"date": "2025 October"
},
"FLUX.2 Dev sdnq-svd-uint4": {
"path": "Disty0/FLUX.2-dev-SDNQ-uint4-svd-r32",
"preview": "Disty0--FLUX.2-dev-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of black-forest-labs/FLUX.2-dev using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"size": 31.89,
"date": "2025 November"
},
"Z-Image-Turbo sdnq-svd-uint4": {
"path": "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
"preview": "Disty0--Z-Image-Turbo-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Tongyi-MAI/Z-Image-Turbo using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 9",
"size": 6.05,
"date": "2025 November"
},
"Qwen-Image-Edit-2511 sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-Edit-2511-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-Edit-2511-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-Edit-2511 using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"date": "2025 December",
"size": 16.09
},
"Qwen-Image-Layered sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-Layered-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-Layered-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-Layered using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"date": "2025 December",
"size": 16.09
},
"Black Forest Labs FLUX.2 Klein 4B sdnq-uint4-dynamic": {
"path": "Disty0/FLUX.2-klein-4B-SDNQ-4bit-dynamic",
"preview": "Disty0--FLUX.2-klein-4B-SDNQ-4bit-dynamic.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-4B using SDNQ.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 5.09,
"date": "2026 January"
},
"Black Forest Labs FLUX.2 Klein 9B sdnq-uint4-dynamic-svd": {
"path": "Disty0/FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32",
"preview": "Disty0--FLUX.2-klein-9B-SDNQ-4bit-dynamic-svd-r32.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-9B using SDNQ with SVD rank 32.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 11.73,
"date": "2026 January"
},
"Qwen-Image-2512 sdnq-svd-uint4": {
"path": "Disty0/Qwen-Image-2512-SDNQ-uint4-svd-r32",
"preview": "Disty0--Qwen-Image-2512-SDNQ-uint4-svd-r32.jpg",
"desc": "Quantization of Qwen/Qwen-Image-2512 using SDNQ: sdnq-svd 4-bit uint with svd rank 32",
"size": 16.09,
"date": "2026 January"
},
"ZAI GLM-Image sdnq-dynamic-uint4": {
"path": "Disty0/GLM-Image-SDNQ-4bit-dynamic",
"preview": "zai-org--GLM-Image.jpg",
@@ -260,6 +218,22 @@
"size": 17.3,
"date": "2026 June"
},
"Krea 2 Base sdnq-hadamard-uint4": {
"path": "vladmandic/Krea-2-Base-sdnq-hadamard-uint4",
"preview": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
"desc": "Krea 2 (K2) Base is the undistilled foundation model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. The base checkpoint is intended for fine-tuning and LoRA training; LoRAs trained on it apply to Krea 2 Turbo.",
"extras": "sampler: Default, cfg_scale: 4.5, steps: 52",
"size": 10.3,
"date": "2026 June"
},
"Black Forest Labs FLUX.2 Klein 9B KV sdnq-uint4-dynamic-svd": {
"path": "vladmandic/Flux.2-Klein-9B-KV-sdnq-hadamard-uint4",
"preview": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
"desc": "Dynamic 4-bit quantization of black-forest-labs/FLUX.2-klein-9B-KV using SDNQ with Hadamard.",
"extras": "sampler: Default, cfg_scale: 1.0, steps: 4",
"size": 12.26,
"date": "2026 July"
},
"Krea 2 Turbo sdnq-hadamard-uint4": {
"path": "vladmandic/Krea-2-Turbo-sdnq-hadamard-uint4",
"preview": "CalamitousFelicitousness--Krea-2-Turbo-Diffusers.jpg",
@@ -268,12 +242,72 @@
"size": 10.54,
"date": "2026 July"
},
"Krea 2 Base sdnq-hadamard-uint4": {
"path": "vladmandic/Krea-2-Base-sdnq-hadamard-uint4",
"preview": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
"desc": "Krea 2 (K2) Base is the undistilled foundation model of the Krea 2 family, trained from scratch by Krea. A 12.9B-parameter single-stream flow-matching DiT that uses a Qwen3-VL-4B vision-language model as its text encoder and the Qwen-Image VAE. The base checkpoint is intended for fine-tuning and LoRA training; LoRAs trained on it apply to Krea 2 Turbo.",
"extras": "sampler: Default, cfg_scale: 4.5, steps: 52",
"size": 10.3,
"date": "2026 June"
"MiniMaxAI MiniMax-H3 sdnq-uint4": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_4bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 64.80,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 sdnq-uint4 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_4bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 64.80,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint4": {
"path": "OzzyGT/MiniMax_H3_sdnq_4bit_pruned",
"preview": "OzzyGT--MiniMax_H3_sdnq_4bit_pruned.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 23.70,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint4 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_4bit_pruned",
"preview": "OzzyGT--MiniMax_H3_sdnq_4bit_pruned.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 4-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 23.70,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 sdnq-uint8": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_8bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 8-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 32.29,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 sdnq-uint8 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_dynamic_8bit",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 8-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 32.29,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint8": {
"path": "OzzyGT/MiniMax_H3_sdnq_8bit_pruned",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 8-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 32.29,
"date": "2026 August"
},
"MiniMaxAI MiniMax-H3 Pruned sdnq-uint8 Ref2VA": {
"path": "OzzyGT/MiniMax_H3_sdnq_8bit_pruned",
"preview": "MiniMaxAI--MiniMax-H3.jpg",
"subfolder": "ref2va",
"desc": "Quantization of MiniMaxAI/MiniMax-H3 using SDNQ: dynamic 8-bit uint. Video with synchronized audio; in image tabs the model runs in experimental still mode.",
"extras": "sampler: Default",
"size": 32.29,
"date": "2026 August"
}
}
extensions-builtin/sdnext-kanvas
+1 -1
Submodule extensions-builtin/sdnext-kanvas updated: b5fb2c580f...e557439043
extensions-builtin/sdnext-modernui
+1 -1
Submodule extensions-builtin/sdnext-modernui updated: 70d6dcae3e...0bacb152d2
extensions-builtin/sdnq
+1 -1
Submodule extensions-builtin/sdnq updated: ac7b5ecf48...b9b50e4a59
installer.py
+68 -55
View File
@@ -1,5 +1,6 @@
from functools import lru_cache
import os
import re
import sys
import json
import time
@@ -57,6 +58,7 @@ args = Dot({
'use_ipex': False,
'use_cuda': False,
'use_rocm': False,
'use_openvino': False,
'experimental': False,
'test': False,
'tls_selfsign': False,
@@ -383,12 +385,11 @@ def git(arg: str, folder: str | None= None, ignore: bool = False, optional: bool
# reattach as needed as head can get detached
def branch(folder=None):
# if args.experimental:
# return None
t_start = time.time()
if not os.path.exists(os.path.join(folder or os.curdir, '.git')):
return None
branches = []
detached = False
try:
b = git('branch --show-current', folder, optional=True)
if b == '':
@@ -397,20 +398,30 @@ def branch(folder=None):
if len(branches) > 0 and len(marked) > 0:
b = marked[0]
if ('detached' in b or 'HEAD' in b) and len(branches) > 1:
detached = True
b = branches[1].strip()
log.debug(f'Git detached head detected: folder="{folder}" reattach={b}')
log.debug(f'Submodule: folder="{folder}" reattach={b} git detached head detected')
except Exception:
b = git('git rev-parse --abbrev-ref HEAD', folder, optional=True)
if args.experimental or args.skip_git or args.skip_all:
return b
if 'main' in b:
b = 'main'
tgt = 'main'
elif 'master' in b:
b = 'master'
tgt = 'master'
else:
b = b.split('\n')[0].replace('*', '').strip()
log.debug(f'Git submodule: {folder} / {b}')
git(f'checkout {b}', folder, ignore=True, optional=True)
tgt = b.split('\n')[0].replace('*', '').strip()
if (tgt != b) or detached:
log.debug(f'Submodule: folder="{folder}" branch="{b}" target="{tgt}"')
git(f'checkout {tgt}', folder, ignore=True, optional=True)
git('fetch', folder, ignore=True)
git(f'merge --ff-only origin/{tgt}', folder, ignore=True)
else:
log.debug(f'Submodule: folder="{folder}" branch="{b}"')
ts('branch', t_start)
return b
return tgt
# restart process
@@ -552,33 +563,19 @@ def check_python(supported_minors=None, experimental_minors=None, reason=None):
# register sdnq package from github submodule
def register_sdnq(skip=False, devices=None, shared=None):
if not skip:
t_start = time.time()
fn = os.path.join('extensions-builtin', 'sdnq', 'src', 'sdnq', '__init__.py')
name = "sdnq"
spec = importlib.util.spec_from_file_location(name, fn)
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module) # this is where actual import happens
import sdnq # pylint: disable=unused-import # test import
ts('sdnq', t_start)
if devices is not None:
import sdnq
sdnq.sdnext.devices = devices
sdnq.quantizer.devices = devices
sdnq.dequantizer.devices = devices
sdnq.quant_utils.devices = devices
sdnq.kernel_wrappers.devices = devices
if shared is not None:
import sdnq
sdnq.sdnext.shared = shared
sdnq.quantizer.shared = shared
sdnq.dequantizer.shared = shared
sdnq.quant_utils.shared = shared
sdnq.kernel_wrappers.shared = shared
sdnq.common.shared = shared
sdnq.loader.shared = shared
def register_sdnq():
t_start = time.time()
os.environ.setdefault('SDNQ_LOGGER_NAME', 'sd')
if not args.use_openvino:
os.environ.setdefault('SDNQ_USE_OPENVINO_MM', '0')
fn = os.path.join('extensions-builtin', 'sdnq', 'src', 'sdnq', '__init__.py')
name = "sdnq"
spec = importlib.util.spec_from_file_location(name, fn)
module = importlib.util.module_from_spec(spec)
sys.modules[name] = module
spec.loader.exec_module(module) # this is where actual import happens
import sdnq # pylint: disable=unused-import # test import
ts('sdnq', t_start)
# check diffusers version
@@ -653,6 +650,26 @@ def check_onnx():
ts('onnx', t_start)
# check numpy version
def check_numpy():
t_start = time.time()
if args.skip_all or args.skip_requirements:
return
torch_ver = package_version('torch') or ''
ver_match = re.match(r'^(\d+)\.(\d+)', torch_ver)
if ver_match:
torch_major, torch_minor = map(int, ver_match.groups())
else:
torch_major, torch_minor = 0, 0
if (torch_major, torch_minor) < (2, 11):
install('numpy==2.1.2', 'numpy', ignore=True)
install('scipy==1.14.1', 'scipy', ignore=True)
else:
install('numpy==2.4.6', 'numpy', ignore=True)
install('scipy==1.18.1', 'scipy', ignore=True)
ts('numpy', t_start)
def install_cuda():
t_start = time.time()
log.info('CUDA: nVidia toolkit detected')
@@ -660,7 +677,7 @@ def install_cuda():
if args.use_nightly:
cmd = os.environ.get('TORCH_COMMAND', '--upgrade --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cu132 --extra-index-url https://download.pytorch.org/whl/nightly/cu130')
else:
cmd = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+cu132 torchvision==0.28.0+cu132 --index-url https://download.pytorch.org/whl/cu132')
cmd = os.environ.get('TORCH_COMMAND', 'torch==2.14.0+cu132 torchvision==0.29.0+cu132 --index-url https://download.pytorch.org/whl/cu132')
return cmd
@@ -732,7 +749,7 @@ def install_rocm_zluda():
zluda_installer.load()
except Exception as e:
log.error(f'Load ZLUDA: {e}')
else: # TODO rocm: switch to pytorch source when it becomes available
else:
if device is None:
log.error('ROCm: no agent found - make sure that graphics driver is installed and up to date')
if device is not None and device.therock is not None:
@@ -821,10 +838,10 @@ def install_openvino():
if sys.platform == 'darwin':
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.11.0 torchvision==0.26.0')
else:
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.11.0+cpu torchvision==0.26.0 --index-url https://download.pytorch.org/whl/cpu')
torch_command = os.environ.get('TORCH_COMMAND', 'torch==2.13.0+cpu torchvision==0.28.0 --index-url https://download.pytorch.org/whl/cpu')
if not (args.skip_all or args.skip_requirements):
install(os.environ.get('OPENVINO_COMMAND', 'openvino==2026.2.1'), 'openvino')
install(os.environ.get('OPENVINO_COMMAND', 'openvino==2026.3.1'), 'openvino')
ts('openvino', t_start)
return torch_command
@@ -1132,7 +1149,7 @@ def list_extensions_folder(folder, quiet=False):
disabled_extensions = opts.get('disabled_extensions', [])
enabled_extensions = [x for x in os.listdir(folder) if os.path.isdir(os.path.join(folder, x)) and x not in disabled_extensions and not x.startswith('.')]
if not quiet:
log.info(f'Extensions: path="{folder}" enabled={enabled_extensions}')
log.info(f'Extensions: path="{folder}" available={enabled_extensions}')
return enabled_extensions
@@ -1289,13 +1306,6 @@ def install_pydantic():
reload('pydantic', '2.13.4')
def install_scipy():
if args.new or (sys.version_info >= (3, 14)):
install('scipy==1.17.1', ignore=True, quiet=True)
else:
install('scipy==1.14.1', ignore=True, quiet=True)
def install_opencv():
install('opencv-python==4.13.0.92', ignore=True, quiet=True)
install('opencv-python-headless==4.13.0.92', ignore=True, quiet=True)
@@ -1324,7 +1334,6 @@ def install_insightface():
def install_optional():
t_start = time.time()
log.info('Installing optional requirements...')
install('pillow-heif')
install('addict')
install('yapf')
install('--no-build-isolation git+https://github.com/Disty0/BasicSR@23c1fb6f5c559ef5ce7ad657f2fa56e41b121754', 'basicsr', ignore=True, quiet=True)
@@ -1335,11 +1344,14 @@ def install_optional():
install('Cython', ignore=True, quiet=True)
install('gguf', ignore=True, quiet=True)
install('hf_transfer', ignore=True, quiet=True)
install('hf_xet', ignore=True, quiet=True)
install('nvidia-ml-py', ignore=True, quiet=True)
install('pillow-heif')
install('pillow-jxl-plugin==1.3.7', ignore=True, quiet=True)
install('ultralytics==8.4.67', ignore=True, quiet=True)
install('open-clip-torch', no_deps=True, quiet=True)
install('runai_model_streamer', ignore=True, quiet=True)
install('facexlib', ignore=True, quiet=True)
install('omegaconf', ignore=True, quiet=True)
install('git+https://github.com/tencent-ailab/IP-Adapter.git', 'ip_adapter', ignore=True, quiet=True)
# install('git+https://github.com/openai/CLIP.git', 'clip', quiet=True, no_build_isolation=True)
ts('optional', t_start)
@@ -1360,7 +1372,6 @@ def install_requirements():
log.info('Install requirements: this may take a while...')
pip('install -r requirements.txt')
if args.optional:
quick_allowed = False
install_optional()
log.info('Install: verifying requirements')
if args.new:
@@ -1375,7 +1386,6 @@ def install_requirements():
install_compel()
install_pydantic()
install_opencv()
install_scipy()
if args.profile:
pr.disable()
print_profile(pr, 'Requirements')
@@ -1384,7 +1394,10 @@ def install_requirements():
# set environment variables controlling the behavior of various libraries
def set_environment():
log.debug('Setting environment tuning')
from modules.logger import console
log.debug(f'Console: terminal={console.is_terminal} width={console.width} height={console.height} color={console.color_system} legacy={console.legacy_windows}')
os.environ.setdefault('ACCELERATE', 'True')
os.environ.setdefault('ATTN_PRECISION', 'fp16')
os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100')
@@ -1540,7 +1553,7 @@ def check_ui(ver):
return
t_start = time.time()
if not same(ver):
log.debug(f'Branch mismatch: {ver}')
log.debug(f'Branch mismatch: module=ModernUI {ver}')
try:
if 'dev' in ver['branch']:
target = 'dev'
@@ -1567,7 +1580,7 @@ def check_kanvas(ver):
return
t_start = time.time()
if not same(ver):
log.debug(f'Branch mismatch: {ver}')
log.debug(f'Branch mismatch: module=Kanvas {ver}')
try:
if 'dev' in ver['branch']:
target = 'dev'
@@ -1702,7 +1715,7 @@ def check_version(reset=True): # pylint: disable=unused-argument
else:
dt = commits["commit"]["commit"]["author"]["date"]
commit = commits["commit"]["sha"][:8]
log.info(f'Version: app=sd.next latest={dt} hash={commit} branch={branch_name}')
log.info(f'Version: app="sd.next" latest={dt} hash={commit} branch={branch_name}')
except Exception as e:
log.error(f'Repository failed to check version: {e} {commits}')
ts('latest', t_start)
launch.py
+2 -1
View File
@@ -238,7 +238,7 @@ def main():
init_args() # setup argparser and default folders
installer.args = args
installer.setup_logging(debug=args.debug, trace=args.trace, filename=args.log)
log.info('Starting SD.Next')
log.info('Starting: [bold cyan]SD.Next[/]')
installer.get_logfile()
try:
sys.excepthook = installer.custom_excepthook
@@ -262,6 +262,7 @@ def main():
installer.install('uv', 'uv')
installer.install_gradio()
installer.check_torch()
installer.check_numpy()
installer.check_onnx()
installer.check_transformers()
installer.check_diffusers()
models/Reference/CalamitousFelicitousness--Anima-2.9B-Preview-v1-Diffusers.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 53 KiB

models/Reference/inclusionAI--LLaDA-Image-Turbo.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 52 KiB

models/Reference/inclusionAI--LLaDA-Image.jpg
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 41 KiB

modules/api/api.py
+6 -1
View File
@@ -29,7 +29,7 @@ class Api:
self.router = APIRouter()
if shared.cmd_opts.docs:
docs.create_docs(app)
docs.create_redocs(app)
# docs.create_redocs(app)
self.app = app
self.queue_lock = queue_lock
self.generate = generate.APIGenerate(queue_lock)
@@ -178,6 +178,11 @@ class Api:
from modules.api import upload
upload.register_api()
# rate limiter
from modules.api.validate import init_limiter
init_limiter()
def add_api_route(self, path: str, fn, auth: bool = True, **kwargs):
if auth and self.credentials:
deps = list(kwargs.get('dependencies', []))
modules/api/docs.py
+3 -2
View File
@@ -1,8 +1,6 @@
import json
from starlette.responses import HTMLResponse
from fastapi import FastAPI
from fastapi.openapi.docs import get_redoc_html, swagger_ui_default_parameters
from fastapi.encoders import jsonable_encoder
def get_swagger_ui_html(*,
@@ -16,6 +14,8 @@ def get_swagger_ui_html(*,
init_oauth: dict | None = None,
swagger_ui_parameters: dict | None = None,
) -> HTMLResponse:
from fastapi.encoders import jsonable_encoder
from fastapi.openapi.docs import swagger_ui_default_parameters
current_swagger_ui_parameters = swagger_ui_default_parameters.copy()
if swagger_ui_parameters:
current_swagger_ui_parameters.update(swagger_ui_parameters)
@@ -79,6 +79,7 @@ def create_docs(app: FastAPI):
def create_redocs(app: FastAPI):
from fastapi.openapi.docs import get_redoc_html
@app.get("/redocs", include_in_schema=False) # override for the default fastapi redocs route
async def custom_redoc_html():
res = get_redoc_html(
modules/api/helpers.py
+1
View File
@@ -46,6 +46,7 @@ def decode_base64_to_image(encoding, quiet=False):
decoded = base64.b64decode(encoding)
data = io.BytesIO(decoded)
image = Image.open(data)
image = image.convert('RGB')
return image
except Exception as e:
log.warning(f'API cannot decode image: {e}')
modules/api/models.py
+4 -4
View File
@@ -366,11 +366,11 @@ class ReqPromptEnhance(BaseModel):
repetition_penalty: Optional[float] = Field(title="Repetition penalty", default=None, description="Penalizes repeated tokens to reduce repetition (1.0=no penalty)")
top_k: Optional[int] = Field(title="Top K", default=None, description="Limits token selection to the K most likely candidates")
top_p: Optional[float] = Field(title="Top P", default=None, description="Nucleus sampling threshold (0-1)")
thinking: bool = Field(title="Thinking", default=False, description="Enable thinking/reasoning mode")
keep_thinking: bool = Field(title="Keep thinking", default=False, description="Keep thinking tokens in output")
use_vision: bool = Field(title="Use vision", default=True, description="Use vision if model supports it")
thinking: Optional[bool] = Field(title="Thinking", default=False, description="Enable thinking/reasoning mode")
keep_thinking: Optional[bool] = Field(title="Keep thinking", default=False, description="Keep thinking tokens in output")
use_vision: Optional[bool] = Field(title="Use vision", default=True, description="Use vision if model supports it")
prefill: Optional[str] = Field(title="Prefill", default=None, description="Text to prefill the model response with")
keep_prefill: bool = Field(title="Keep prefill", default=False, description="Keep prefill text in the output")
keep_prefill: Optional[bool] = Field(title="Keep prefill", default=False, description="Keep prefill text in the output")
custom_args: Optional[str] = Field(title="Custom args", default=None, description="Custom arguments for the model")
process_words: Optional[str] = Field(title="Banned words", default=None, description="List of words to process")
semantic_threshold: Optional[float] = Field(title="Semantic threshold", default=None, description="Semantic similarity threshold for processed words")
modules/api/process.py
+2 -1
View File
@@ -226,6 +226,7 @@ class APIProcess:
if len(instance) == 0:
raise HTTPException(status_code=500, detail="Prompt enhancement script not found")
instance = instance[0]
decoded = decode_base64_to_image(req.image) if req.image else None
prompt = instance.enhance(
model=model,
prompt=req.prompt,
@@ -244,7 +245,7 @@ class APIProcess:
use_vision=req.use_vision,
prefill=req.prefill or '',
keep_prefill=req.keep_prefill,
image=decode_base64_to_image(req.image) if req.image else None,
image=decoded,
seed=seed,
nsfw=req.nsfw,
custom_args=req.custom_args,
modules/api/validate.py
+28 -18
View File
@@ -7,35 +7,37 @@ request_cost = {
"/file": 0,
"/internal/progress": 0,
"/run/predict": 0,
"/sdapi/v1/control": 5,
"/sdapi/v1/img2img": 5,
"/sdapi/v1/txt2img": 5,
"/sdapi/v1/video": 5,
"/sdapi/v1/browser/thumb": 0,
"/sdapi/v1/network/thumb": 0,
"/sdapi/v1/txt2img": 5,
"/sdapi/v1/img2img": 5,
"/sdapi/v1/control": 5,
"/sdapi/v1/video": 5,
}
log_cost = {
"/.well-known/appspecific/com.chrome.devtools.json": -1,
"/info": -1,
"/file": -1,
"/token": -1,
"/theme.css": -1,
"/sdapi/v1/browser/thumb": -1,
"/sdapi/v1/network/thumb": -1,
"/run/predict": -1,
"/queue/join": -1,
"/info": -1,
"/icon": -1,
"/internal/progress": -1,
"/sdapi/v1/version": -1,
"/sdapi/v1/log": -1,
"/sdapi/v1/torch": -1,
"/queue/join": -1,
"/run/predict": -1,
"/theme.css": -1,
"/token": -1,
"/sdapi/v1/checkpoint": -1,
"/sdapi/v1/gpu-smi": -1,
"/sdapi/v1/gpu": -1,
"/sdapi/v1/loaded-loras": -1,
"/sdapi/v1/log": -1,
"/sdapi/v1/memory": -1,
"/sdapi/v1/platform": -1,
"/sdapi/v1/checkpoint": -1,
"/sdapi/v1/loaded-loras": -1,
"/sdapi/v1/gpu-smi": -1,
"/sdapi/v1/status": 60,
"/sdapi/v1/progress": 60,
"/sdapi/v1/start": -1,
"/sdapi/v1/status": 60,
"/sdapi/v1/torch": -1,
"/sdapi/v1/version": -1,
"/sdapi/v1/browser/thumb": -1,
"/sdapi/v1/network/thumb": -1,
}
log_exclude_suffix = ['.css', '.js', '.ico', '.svg']
log_exclude_prefix = ['/assets']
@@ -99,6 +101,13 @@ def get_api_stats():
limiter.stats()
def init_limiter():
global limiter # pylint: disable=global-statement
from modules.shared import opts, cmd_opts
if opts.server_rate_limit != limiter.request_limit:
limiter = Limiter(opts.server_rate_limit, cmd_opts.subpath, cmd_opts.profile)
def validate_request(client, endpoint):
global limiter # pylint: disable=global-statement
from modules.shared import opts, cmd_opts
@@ -114,6 +123,7 @@ def validate_request(client, endpoint):
limiter.summary[key] += 1
return limiter.check_request(client, api)
def validate_log(client, endpoint):
api = re.match(r"^[^?#&=]+", endpoint).group(0)
if (limiter.subpath is not None) and (len(limiter.subpath) > 0) and api.startswith(limiter.subpath): # strip subpath from api for logging
modules/api/video.py
+3 -1
View File
@@ -22,8 +22,9 @@ class ReqVideo(BaseModel):
frames: int = Field(default=17, ge=1, le=1024, title="Frames", description="Number of frames; 1 produces a single still image on workflow models")
steps: int = Field(default=50, ge=1, le=200, title="Steps", description="Number of inference steps")
sampler_name: str = Field(default="Default", title="Sampler", description="Sampler name; Default keeps the model scheduler")
sampler_shift: float = Field(default=-1.0, title="Sampler shift", description="Scheduler flow shift; -1 keeps the model default")
sampler_shift: float = Field(default=-1.0, title="Sampler shift", description="Scheduler flow shift, the video schedule on models with a separate audio schedule; -1 keeps the model default")
dynamic_shift: bool = Field(default=False, title="Dynamic shift", description="Enable dynamic scheduler shifting")
audio_shift: float = Field(default=-1.0, title="Audio shift", description="Audio schedule shift on models with a separate audio scheduler; -1 keeps the model default")
seed: int = Field(default=-1, title="Seed", description="Generation seed; -1 for random")
guidance_scale: float = Field(default=-1.0, title="Guidance scale", description="CFG scale; -1 keeps the model default")
guidance_true: float = Field(default=-1.0, title="True guidance", description="True CFG scale; -1 keeps the model default")
@@ -173,6 +174,7 @@ class APIVideo:
sampler_name=sampler_name,
sampler_shift=req.sampler_shift,
dynamic_shift=req.dynamic_shift,
audio_shift=req.audio_shift,
seed=req.seed,
guidance_scale=req.guidance_scale,
guidance_true=req.guidance_true,
modules/attention.py
-355
View File
@@ -1,355 +0,0 @@
from functools import wraps
import torch
from modules import rocm, errors, devices
from modules.logger import log
from installer import install, installed, torch_info
def set_dynamic_attention():
try:
sdpa_pre_dyanmic_atten = torch.nn.functional.scaled_dot_product_attention
from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
torch.nn.functional.scaled_dot_product_attention = dynamic_scaled_dot_product_attention
torch_info.set(attention='dynamic')
return sdpa_pre_dyanmic_atten
except Exception as err:
log.error(f'Torch attention: type="dynamic attention" {err}')
return None
def set_sdnq_attention():
try:
from modules import shared
from sdnq.kernels.triton_atten import sdnq_triton_atten
sdpa_pre_sdnq_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sdnq_atten)
def sdpa_sdnq_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
if (
query.device.type != "cpu"
and (query.shape[-2] >= 32 and key.shape[-2] >= 32)
and (query.shape[-2] > 512 or key.shape[-2] > 512) # Skip TE
and query.shape[-3] > 1 # Skip VAE
):
return sdnq_triton_atten(
query=query, key=key, value=value, attn_mask=attn_mask,
is_causal=is_causal, scale=scale, enable_gqa=enable_gqa,
matmul_dtype=shared.opts.sdnq_attention_matmul_type,
pv_matmul_dtype=shared.opts.sdnq_attention_pv_matmul_type,
smooth_k=shared.opts.sdnq_attention_smooth_k,
use_hadamard=shared.opts.sdnq_attention_use_hadamard,
hadamard_group_size=shared.opts.sdnq_attention_hadamard_group_size,
quantize_fp32=shared.opts.sdnq_attention_quantize_fp32,
use_fp16_accum=shared.opts.sdnq_attention_use_fp16_accum,
)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_sdnq_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_sdnq_atten
torch_info.set(attention='sdnq')
log.debug(f'Torch attention: type="SDNQ attention" matmul={shared.opts.sdnq_attention_matmul_type}:{shared.opts.sdnq_attention_pv_matmul_type} smooth={shared.opts.sdnq_attention_smooth_k} hadamard={shared.opts.sdnq_attention_use_hadamard} fp16_accum={shared.opts.sdnq_attention_use_fp16_accum}')
except Exception as err:
log.error(f'Torch attention: type="SDNQ attention" {err}')
def set_triton_flash_attention(backend: str):
try:
if backend in {"rocm", "zluda"}: # flash_attn_triton_amd only works with AMD
from modules.flash_attn_triton_amd import interface_fa
sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_triton_flash_atten)
def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_triton = (
query.shape[-1] <= 128
and attn_mask is None
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_triton:
if scale is None:
scale = query.shape[-1] ** (-0.5)
head_size_og = query.size(3)
if head_size_og % 8 != 0:
query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
out_padded = torch.zeros_like(query)
interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
return out_padded[..., :head_size_og].transpose(1, 2)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten
torch_info.set(attention='triton')
log.debug('Torch attention: type="Triton Flash attention"')
except Exception as err:
log.error(f'Torch attention: type="Triton Flash attention" {err}')
def set_flex_attention():
try:
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
def flex_attention_causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
return q_idx >= kv_idx
sdpa_pre_flex_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_flex_atten)
def sdpa_flex_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor: # pylint: disable=unused-argument
score_mod = None
block_mask = None
if attn_mask is not None:
batch_size, num_heads = query.shape[:2]
seq_len_q = query.shape[-2]
seq_len_kv = key.shape[-2]
if attn_mask.ndim == 2:
attn_mask = attn_mask.view(attn_mask.shape[0], 1, attn_mask.size[1], 1)
attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv)
if attn_mask.dtype == torch.bool:
def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
block_mask = create_block_mask(mask_mod, batch_size, None, seq_len_q, seq_len_kv, device=query.device)
else:
def score_mod_fn(score, batch_idx, head_idx, q_idx, kv_idx):
return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
score_mod = score_mod_fn
elif is_causal:
block_mask = create_block_mask(flex_attention_causal_mask, query.shape[0], query.shape[1], query.shape[-2], key.shape[-2], device=query.device)
return flex_attention(query, key, value, score_mod=score_mod, block_mask=block_mask, scale=scale, enable_gqa=enable_gqa)
torch.nn.functional.scaled_dot_product_attention = sdpa_flex_atten
torch_info.set(attention="flex")
log.debug('Torch attention: type="Flex attention"')
except Exception as err:
log.error(f'Torch attention: type="Flex attention" {err}')
def set_ck_flash_attention(backend: str, device: torch.device):
try:
if backend == "rocm":
if not installed('flash-attn'):
log.info('Torch attention: type="Flash attention" building...')
agent = rocm.Agent(device)
install(rocm.get_flash_attention_command(agent), reinstall=True)
else:
install('flash-attn')
from flash_attn import flash_attn_func
sdpa_pre_flash_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_flash_atten)
def sdpa_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_flash = (
query.shape[-1] <= 128
and attn_mask is None
and query.dtype != torch.float32
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_flash:
is_unsqueezed = False
if query.dim() == 3:
query = query.unsqueeze(0)
is_unsqueezed = True
if key.dim() == 3:
key = key.unsqueeze(0)
if value.dim() == 3:
value = value.unsqueeze(0)
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
if is_unsqueezed:
attn_output = attn_output.squeeze(0)
return attn_output
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten
torch_info.set(attention="flash")
log.debug('Torch attention: type="Flash attention"')
except Exception as err:
log.error(f'Torch attention: type="Flash attention" {err}')
def set_sage_attention(backend: str, device: torch.device):
try:
install('sageattention')
use_cuda_backend = False
if (backend == "cuda") and (torch.cuda.get_device_capability(device) == (8, 6)):
use_cuda_backend = True # Detect GPU architecture - sm86 confirmed to need CUDA backend workaround as Sage Attention + Triton causes NaNs
try:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
except Exception:
use_cuda_backend = False
if use_cuda_backend:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn_qk_int8_pv_fp16_cuda(
q=query, k=key, v=value,
tensor_layout="HND",
is_causal=is_causal,
sm_scale=scale,
return_lse=False,
pv_accum_dtype="fp32",
)
else:
from sageattention import sageattn
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn(
q=query, k=key, v=value,
attn_mask=None,
dropout_p=0.0,
is_causal=is_causal,
scale=scale,
)
sdpa_pre_sage_atten = torch.nn.functional.scaled_dot_product_attention
@wraps(sdpa_pre_sage_atten)
def sdpa_sage_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: torch.Tensor | None = None, dropout_p: float = 0.0, is_causal: bool = False, scale: float | None = None, enable_gqa: bool = False, **kwargs) -> torch.Tensor:
use_sage = (
query.shape[-1] in {128, 96, 64}
and attn_mask is None
and query.device.type != "cpu"
and key.device == query.device
and value.device == query.device
)
if use_sage:
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
# Call preselected sage attention implementation
return sage_attn_impl(query, key, value, is_causal, scale)
else:
if enable_gqa:
kwargs["enable_gqa"] = enable_gqa
return sdpa_pre_sage_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
torch.nn.functional.scaled_dot_product_attention = sdpa_sage_atten
torch_info.set(attention="sage")
log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
except Exception as err:
log.error(f'Torch attention: type="Sage attention" {err}')
def set_diffusers_attention(pipe, quiet = False):
from modules import shared
import diffusers.models.attention_processor as p
def set_attn(pipe, attention, name: str | None = None):
if attention is None:
return
# other models uses their own attention processor
if getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"):
try:
pipe.unet.set_attn_processor(attention)
except Exception as e:
if 'Nunchaku' in pipe.unet.__class__.__name__:
pass
else:
log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}')
log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
if shared.opts.cross_attention_optimization == "Disabled":
torch_info.set(attention="disabled")
elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
devices.set_sdpa_params()
# set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product")
elif shared.opts.cross_attention_optimization == "xFormers":
if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
torch_info.set(attention="xformers")
pipe.enable_xformers_memory_efficient_attention()
else:
log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
elif shared.opts.cross_attention_optimization == "Batch matrix-matrix":
torch_info.set(attention="bmm")
set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix")
elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
torch_info.set(attention="dynamic_bmm")
set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM")
if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"):
if shared.opts.attention_slicing:
pipe.enable_attention_slicing()
else:
pipe.disable_attention_slicing()
log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}")
pipe.current_attn_name = shared.opts.cross_attention_optimization
orig_get_kernel = None
def get_kernel_hijack(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code: bool | list[str] = False): # pylint: disable=unused-argument
log.debug(f'Attention dispatcher hub: repo="{repo_id}" revision={revision} version={version} backend={backend}')
user_agent = 'kernels/0.16.0'
module = None
try:
module = orig_get_kernel(repo_id, revision=revision, version=version, backend=backend, user_agent=user_agent, trust_remote_code=True)
except Exception as e:
log.error(f'Attention dispatcher hub: {e}')
errors.display(e, 'kernels')
return module
def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument
from huggingface_hub import HfApi
return HfApi(library_name="kernels", user_agent="donottrack")
def hijack_kernels():
global orig_get_kernel # pylint: disable=global-statement
try:
install('kernels==0.16.0')
import kernels
import kernels.utils
log.debug(f'Attention dispatcher: kernels={kernels.__version__}')
if orig_get_kernel is None:
orig_get_kernel = kernels.get_kernel
kernels.get_kernel = get_kernel_hijack
kernels.utils._get_hf_api = get_hf_api_hijack # pylint: disable=protected-access
from diffusers.utils import import_utils
import_utils._kernels_available = True # pylint: disable=protected-access
import_utils._kernels_version = kernels.__version__ # pylint: disable=protected-access
except Exception as e:
log.error(f'Attention dispatcher kernels: {e}')
return
def set_attention_dispatcher(pipe):
from modules import shared
attn = shared.opts.hf_attention.strip().lower()
if pipe is None or not hasattr(pipe, 'transformer') or not hasattr(pipe.transformer, 'set_attention_backend'):
return
from diffusers.models import attention_dispatch as a
backends = [b.value for b in a._AttentionBackendRegistry.list_backends()] # pylint: disable=protected-access
# https://huggingface.co/docs/kernels/index
# https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends
if 'hub' in attn:
hijack_kernels()
prev = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
if attn in backends:
try:
pipe.transformer.set_attention_backend(attn)
except Exception as e:
log.error(f'Attention dispatcher: target={attn} {e}')
current = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
log.debug(f'Attention dispatcher: target={attn} previous={prev[0].value} active={current[0]} list={backends}')
elif len(attn) > 0:
log.warning(f'Attention dispatcher: active={prev[0].value} list={backends} target={attn} not found')
else:
log.debug(f'Attention dispatcher: active={prev[0].value} list={backends}')
modules/attention/__init__.py
+12
View File
@@ -0,0 +1,12 @@
"""Attention backends: one scaled_dot_product_attention router over the registered backends, the per-generation context, and the diffusers-side processor and dispatcher setup."""
from modules.attention.registry import AttentionBackend, AttentionCall, Constraints, Platform, Registry, registry
from modules.attention.router import Plan, PlanEntry, build_plan, get_plan, install_router, reapply, reapply_options, report
from modules.attention.dispatcher import set_diffusers_attention, set_attention_dispatcher, list_dispatcher_backends, hijack_kernels, get_kernel_hijack, get_hf_api_hijack
from modules.attention import backends, context, debug
__all__ = [
'AttentionBackend', 'AttentionCall', 'Constraints', 'Platform', 'Registry', 'registry',
'Plan', 'PlanEntry', 'build_plan', 'get_plan', 'install_router', 'reapply', 'reapply_options', 'report',
'set_diffusers_attention', 'set_attention_dispatcher', 'list_dispatcher_backends', 'hijack_kernels', 'get_kernel_hijack', 'get_hf_api_hijack',
'backends', 'context', 'debug',
]
modules/attention/backends/__init__.py
+10
View File
@@ -0,0 +1,10 @@
"""Built-in backends, registered in ascending priority."""
from modules.attention.registry import registry
from modules.attention.backends import dynamic, flex, triton_amd, flash_ck, sage, sdnq
registry.register(dynamic.backend)
registry.register(flex.backend)
registry.register(triton_amd.backend)
registry.register(flash_ck.backend)
registry.register(sage.backend)
registry.register(sdnq.backend)
modules/attention/backends/dynamic.py
+11
View File
@@ -0,0 +1,11 @@
from modules.attention.registry import AttentionBackend, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from modules import devices
devices.sdpa_pre_dyanmic_atten = original # the sliced path calls this pin for every slice
from modules.sd_hijack_dynamic_atten import dynamic_scaled_dot_product_attention
return dynamic_scaled_dot_product_attention
backend = AttentionBackend(name='dynamic', label='Dynamic attention', priority=10, prepare=prepare, terminal=True)
modules/attention/backends/flash_ck.py
+48
View File
@@ -0,0 +1,48 @@
from installer import install, installed
from modules import rocm
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
try:
import flash_attn # pylint: disable=unused-import
except ImportError:
log.warning('Attention: type="Flash attention" not installed: starting build, this may take a while...')
if platform.backend == 'rocm':
if not installed('flash-attn'):
log.info('Attention: type="Flash attention" building...')
agent = rocm.Agent(platform.device)
install(rocm.get_flash_attention_command(agent), reinstall=True)
else:
install('--no-build-isolation flash-attn')
from flash_attn import flash_attn_func
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
is_unsqueezed = False
if query.dim() == 3:
query = query.unsqueeze(0)
is_unsqueezed = True
if key.dim() == 3:
key = key.unsqueeze(0)
if value.dim() == 3:
value = value.unsqueeze(0)
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
attn_output = flash_attn_func(q=query, k=key, v=value, dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
if is_unsqueezed:
attn_output = attn_output.squeeze(0)
return attn_output
log.debug('Attention: type="Flash attention"')
return call
backend = AttentionBackend(
name='flash', label='Flash attention', priority=40, prepare=prepare,
constraints=Constraints(max_head_dim=128, allow_mask=False, allow_float32=False, same_device=True),
)
modules/attention/backends/flex.py
+46
View File
@@ -0,0 +1,46 @@
import torch
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from torch.nn.attention.flex_attention import create_block_mask
from modules.attention.sparse import flex as sparse_flex
def causal_mask(b, h, q_idx, kv_idx): # pylint: disable=unused-argument
return q_idx >= kv_idx
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
if selection is not None:
return sparse_flex.attend(query, key, value, selection, scale=scale, enable_gqa=enable_gqa)
# compiled, always: eager flex_attention materializes the whole score matrix, which is
# tens of gigabytes at video sequence lengths and fails in the driver rather than cleanly
flex_attention = sparse_flex.flex_call()
score_mod = None
block_mask = None
if attn_mask is not None:
batch_size, num_heads = query.shape[:2]
seq_len_q = query.shape[-2]
seq_len_kv = key.shape[-2]
attn_mask = attn_mask.expand(batch_size, num_heads, seq_len_q, seq_len_kv) # sdpa masks broadcast over the trailing dims
if attn_mask.dtype == torch.bool:
def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
return attn_mask[batch_idx, head_idx, q_idx, kv_idx]
block_mask = create_block_mask(mask_mod, batch_size, None, seq_len_q, seq_len_kv, device=query.device)
else:
def score_mod_fn(score, batch_idx, head_idx, q_idx, kv_idx):
return score + attn_mask[batch_idx, head_idx, q_idx, kv_idx]
score_mod = score_mod_fn
elif is_causal:
block_mask = create_block_mask(causal_mask, query.shape[0], query.shape[1], query.shape[-2], key.shape[-2], device=query.device)
return flex_attention(query, key, value, score_mod=score_mod, block_mask=block_mask, scale=scale, enable_gqa=enable_gqa)
log.debug('Attention: type="Flex attention"')
return call
backend = AttentionBackend(
name='flex', label='Flex attention', priority=20, prepare=prepare,
constraints=Constraints(min_ndim=4, same_device=True), # flex_attention takes 4d tensors on one device and compiles on cpu
caps=frozenset({'block_mask'}),
)
modules/attention/backends/sage.py
+55
View File
@@ -0,0 +1,55 @@
import torch
from installer import install, installed
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
if not installed('sageattention'):
log.warning('Attention: type="Sage attention" not installed: starting build, this may take a while...')
install('--no-build-isolation git+http://github.com/thu-ml/SageAttention.git', 'sageattention')
use_cuda_backend = False
if platform.backend == 'cuda' and torch.cuda.get_device_capability(platform.device) == (8, 6):
use_cuda_backend = True # sm86 needs the cuda backend, sage attention over triton produces NaNs there
try:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
except Exception:
use_cuda_backend = False
if use_cuda_backend:
from sageattention import sageattn_qk_int8_pv_fp16_cuda
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn_qk_int8_pv_fp16_cuda(
q=query, k=key, v=value,
tensor_layout="HND",
is_causal=is_causal,
sm_scale=scale,
return_lse=False,
pv_accum_dtype="fp32",
)
else:
from sageattention import sageattn
def sage_attn_impl(query, key, value, is_causal, scale):
return sageattn(
q=query, k=key, v=value,
attn_mask=None,
dropout_p=0.0,
is_causal=is_causal,
scale=scale,
)
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
if enable_gqa:
key = key.repeat_interleave(query.size(-3)//key.size(-3), -3)
value = value.repeat_interleave(query.size(-3)//value.size(-3), -3)
return sage_attn_impl(query, key, value, is_causal, scale)
log.debug(f'Attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}')
return call
backend = AttentionBackend(
name='sage', label='Sage attention', priority=50, prepare=prepare,
constraints=Constraints(head_dims=frozenset({64, 96, 128}), allow_mask=False, same_device=True),
)
modules/attention/backends/sdnq.py
+45
View File
@@ -0,0 +1,45 @@
import inspect
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def supports_block_mask(entry) -> bool:
"""Whether the installed sdnq takes a block mask; the chain must not promise what the kernel cannot do."""
try:
return 'block_mask' in inspect.signature(inspect.unwrap(entry)).parameters
except (TypeError, ValueError):
return False
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from modules import shared
from sdnq.kernels.triton_atten import sdnq_triton_atten
options = {
'matmul_dtype': shared.opts.sdnq_attention_matmul_type,
'pv_matmul_dtype': shared.opts.sdnq_attention_pv_matmul_type,
'smooth_k': shared.opts.sdnq_attention_smooth_k,
'use_hadamard': shared.opts.sdnq_attention_use_hadamard,
'hadamard_group_size': shared.opts.sdnq_attention_hadamard_group_size,
'quantize_fp32': shared.opts.sdnq_attention_quantize_fp32,
'use_fp16_accum': shared.opts.sdnq_attention_use_fp16_accum,
}
block_mask = supports_block_mask(sdnq_triton_atten)
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=None): # pylint: disable=unused-argument
if selection is not None:
return sdnq_triton_atten(query=query, key=key, value=value, attn_mask=attn_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, block_mask=selection.keep, block_mask_m=selection.block_q, block_mask_n=selection.block_kv, **options)
return sdnq_triton_atten(query=query, key=key, value=value, attn_mask=attn_mask, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **options)
call.caps = backend.caps if block_mask else frozenset()
if not block_mask and getattr(shared.opts, 'sparse_attention_enabled', False):
log.warning('SDNQ attention: the installed sdnq has no block mask input, sparse attention cannot use it; update the sdnq submodule')
log.debug(f'Attention: type="SDNQ attention" matmul={options["matmul_dtype"]}:{options["pv_matmul_dtype"]} smooth={options["smooth_k"]} hadamard={options["use_hadamard"]} quantize_fp32={options["quantize_fp32"]} fp16_accum={options["use_fp16_accum"]} block_mask={block_mask}')
return call
backend = AttentionBackend(
name='sdnq', label='SDNQ attention', priority=60, prepare=prepare,
constraints=Constraints(min_tokens=32, min_long_side=512, min_heads=2), # sequences of 512 or fewer are text encoders, single-head calls the vae
options=('sdnq_attention_matmul_type', 'sdnq_attention_pv_matmul_type', 'sdnq_attention_smooth_k', 'sdnq_attention_use_hadamard', 'sdnq_attention_hadamard_group_size', 'sdnq_attention_quantize_fp32', 'sdnq_attention_use_fp16_accum'),
caps=frozenset({'block_mask', 'masked_block'}), # the kernel takes attn_mask and block_mask together
)
modules/attention/backends/triton_amd.py
+32
View File
@@ -0,0 +1,32 @@
import torch
from modules.logger import log
from modules.attention.registry import AttentionBackend, Constraints, Platform
def prepare(platform: Platform, original): # pylint: disable=unused-argument
from modules.flash_attn_triton_amd import interface_fa
def call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa): # pylint: disable=unused-argument
if scale is None:
scale = query.shape[-1] ** (-0.5)
head_size_og = query.size(3)
if head_size_og % 8 != 0:
query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8])
key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8])
value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8])
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
out_padded = torch.zeros_like(query)
interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal)
return out_padded[..., :head_size_og].transpose(1, 2)
log.debug('Attention: type="Triton AMD Flash attention"')
return call
backend = AttentionBackend(
name='triton', label='Triton AMD Flash attention', priority=30, prepare=prepare,
constraints=Constraints(max_head_dim=128, allow_mask=False, same_device=True),
platforms=frozenset({'rocm', 'zluda'}),
)
modules/attention/context.py
+111
View File
@@ -0,0 +1,111 @@
"""Per-generation state for attention consumers: the component running, the denoiser forward about to run, and the model."""
from contextlib import contextmanager
from dataclasses import dataclass
import torch
@dataclass
class GenerationContext:
active: bool = False
role: str | None = None # 'transformer', 'te' or 'vae' while a generation runs, None outside one
step: int = 0 # index of the denoiser forward about to run
steps: int = 0 # forwards in the current pass
forwards: int = 0
model_key: tuple[str, str | None] | None = None # pipeline class and denoiser class, for telemetry and the sparse exclusion list
step_buffer: torch.Tensor | None = None # the step as a device scalar updated in place, so compiled readers keep their graph
layout: object | None = None # TokenLayout published by whoever knows the packing, None until something does
current = GenerationContext()
def denoiser_name(pipe) -> str | None:
for name in ('transformer', 'unet'):
module = getattr(pipe, name, None)
if module is not None:
return module.__class__.__name__
return None
# every slot a pipeline can enter once per denoising step, from sd_offload_state.group_offload_main. The aux
# components on that list (decoder, controlnet, prior) are left alone: they pack no attention sequence, and a
# publication from one would clear the layout the denoiser just set
DENOISER_SLOTS = ('transformer', 'unet', 'transformer_2', 'transformer_ref', 'unconditional_transformer')
def install_layout_hook(pipe) -> None:
"""Let a classic pipeline's denoiser publish its own packing: the modular path has its own hook, this is the rest."""
from modules import shared
if pipe is None or not getattr(shared.opts, 'sparse_attention_enabled', False):
return
from modules.attention.sparse import layout as sparse_layout
def publish(denoiser, args, kwargs): # pylint: disable=unused-argument
set_layout(sparse_layout.layout_from_kwargs(kwargs, denoiser.__class__.__name__))
for name in DENOISER_SLOTS:
module = getattr(pipe, name, None)
if module is None or getattr(module, 'sdnext_layout_hook', None) is not None or getattr(module, 'sdnext_state_hook', None) is not None:
continue
module.sdnext_layout_hook = module.register_forward_pre_hook(publish, with_kwargs=True)
def begin(pipe, steps: int = 0) -> None:
from modules import devices
current.active = True
current.role = 'transformer'
current.layout = None
current.model_key = (pipe.__class__.__name__, denoiser_name(pipe)) if pipe is not None else None
install_layout_hook(pipe)
device = devices.device if devices.device is not None else torch.device('cpu')
if current.step_buffer is None or current.step_buffer.device != device:
current.step_buffer = torch.zeros((), dtype=torch.int64, device=device)
new_pass(steps)
def new_pass(steps: int = 0) -> None:
"""Restart the step count for a denoising pass: base, hires or refiner."""
current.steps = int(steps or 0)
current.forwards = 0
set_step(0)
def set_step(step: int) -> None:
current.step = int(step)
if current.step_buffer is not None:
current.step_buffer.fill_(current.step)
def tick(step: int | None = None) -> None:
"""Advance to the next forward: the classic callback passes the completed step plus one, the modular pre-hook passes nothing and counts forwards."""
set_step(current.forwards if step is None else step)
current.forwards = current.step + 1
def set_layout(layout) -> None:
"""Publish what the packed sequence holds; callers that know the packing set this per forward."""
current.layout = layout
def end() -> None:
from modules.attention import debug
current.active = False
current.role = None
current.model_key = None
current.layout = None
new_pass(0)
debug.end_generation()
def set_role(name: str | None) -> None:
current.role = name
@contextmanager
def role(name: str):
previous = current.role
current.role = name
try:
yield
finally:
current.role = previous
modules/attention/debug.py
+42
View File
@@ -0,0 +1,42 @@
"""Opt-in route tracing for the sdpa router, enabled by SD_ATTN_DEBUG."""
import os
import torch
from modules.logger import log
from modules.attention import context
enabled = os.environ.get('SD_ATTN_DEBUG', None) is not None
seen: set[tuple] = set()
counts: dict[tuple, int] = {}
def observe(name: str, query: torch.Tensor, key: torch.Tensor, attn_mask: torch.Tensor | None) -> None:
"""Log each distinct route once: backend, component role, step, shapes, dtype, mask presence and whether the inputs are contiguous; count every call."""
contiguous = query.is_contiguous() and key.is_contiguous()
signature = (name, context.current.role, tuple(query.shape), tuple(key.shape), str(query.dtype), attn_mask is not None, contiguous)
counts[signature] = counts.get(signature, 0) + 1
if signature in seen:
return
seen.add(signature)
log.debug(f'Attention route: backend={name} role={context.current.role} step={context.current.step} q={list(query.shape)} k={list(key.shape)} dtype={query.dtype} mask={attn_mask is not None} contiguous={contiguous}')
def summary() -> list[str]:
"""One line per route with its call count since the last generation, busiest first."""
lines = []
for signature, count in sorted(counts.items(), key=lambda item: -item[1]):
name, role, q_shape, k_shape, dtype, masked, contiguous = signature
lines.append(f'backend={name} role={role} q={list(q_shape)} k={list(k_shape)} dtype={dtype} mask={masked} contiguous={contiguous} calls={count}')
return lines
def end_generation() -> None:
"""Log the route counts of the generation that just ended and start the next count."""
if enabled and counts:
for line in summary():
log.debug(f'Attention routes: {line}')
counts.clear()
def reset() -> None:
seen.clear()
counts.clear()
modules/attention/dispatcher.py
+126
View File
@@ -0,0 +1,126 @@
from modules import errors
from modules.logger import log
from installer import install, torch_info
def set_xformers_attention(pipe):
try:
# install('xformers')
import xformers
log.debug(f'Attention: xFormers={xformers.__version__}')
import diffusers.utils.import_utils
diffusers.utils.import_utils._xformers_available = True # pylint: disable=protected-access
diffusers.utils.import_utils._xformers_version = xformers.__version__ # pylint: disable=protected-access
import diffusers.models.attention_processor
import importlib
importlib.reload(diffusers.models.attention_processor)
# diffusers.models.attention_processor.xformers = xformers
except Exception as e:
log.error(f'Attention: xFormers {e}')
return
if hasattr(pipe, 'enable_xformers_memory_efficient_attention'):
torch_info.set(attention="xformers")
pipe.enable_xformers_memory_efficient_attention()
else:
log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}")
def set_diffusers_attention(pipe, quiet = False):
from modules import shared, attention
log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"')
attention.reapply()
plan = attention.get_plan()
if plan is not None and (plan.entries or plan.terminal):
pass # already set by router
elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers
pass # attention.reapply already called devices.set_sdpa_params
elif shared.opts.cross_attention_optimization == "xFormers":
set_xformers_attention(pipe)
elif shared.opts.cross_attention_optimization == "Disabled" or shared.opts.cross_attention_optimization == "Default":
torch_info.set(attention="default")
else:
log.warning(f'Attention: cls={pipe.__class__.__name__} method="{shared.opts.cross_attention_optimization}" not applied')
if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"):
if shared.opts.attention_slicing == "Enabled":
pipe.enable_attention_slicing()
else:
pipe.disable_attention_slicing()
log.debug(f"Attention: slicing={shared.opts.attention_slicing}")
pipe.current_attn_name = shared.opts.cross_attention_optimization
orig_get_kernel = None
def get_kernel_hijack(repo_id, revision=None, version=None, backend=None, user_agent=None, trust_remote_code: bool | list[str] = False): # pylint: disable=unused-argument
log.debug(f'Attention dispatcher hub: repo="{repo_id}" revision={revision} version={version} backend={backend}')
user_agent = 'kernels/0.16.0'
module = None
try:
module = orig_get_kernel(repo_id, revision=revision, version=version, backend=backend, user_agent=user_agent, trust_remote_code=True)
except Exception as e:
log.error(f'Attention dispatcher hub: {e}')
errors.display(e, 'kernels')
return module
def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument
from huggingface_hub import HfApi
return HfApi(library_name="kernels", user_agent="donottrack")
def hijack_kernels():
global orig_get_kernel # pylint: disable=global-statement
try:
install('kernels==0.16.1')
import kernels
import kernels.utils
log.debug(f'Attention dispatcher: kernels={kernels.__version__}')
if orig_get_kernel is None:
orig_get_kernel = kernels.get_kernel
kernels.get_kernel = get_kernel_hijack
kernels.utils._get_hf_api = get_hf_api_hijack # pylint: disable=protected-access
from diffusers.utils import import_utils
import_utils._kernels_available = True # pylint: disable=protected-access
import_utils._kernels_version = kernels.__version__ # pylint: disable=protected-access
except Exception as e:
log.error(f'Attention dispatcher kernels: {e}')
return
def set_attention_dispatcher(pipe):
from modules import shared
attn = shared.opts.hf_attention.strip().lower()
if pipe is None or not hasattr(pipe, 'transformer') or not hasattr(pipe.transformer, 'set_attention_backend'):
return
from diffusers.models import attention_dispatch as a
backends = [b.value for b in a._AttentionBackendRegistry.list_backends()] # pylint: disable=protected-access
# https://huggingface.co/docs/kernels/index
# https://huggingface.co/docs/diffusers/optimization/attention_backends#available-backends
if 'hub' in attn:
hijack_kernels()
prev = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
if attn in backends:
try:
pipe.transformer.set_attention_backend(attn)
except Exception as e:
log.error(f'Attention dispatcher: target={attn} {e}')
current = a._AttentionBackendRegistry.get_active_backend() # pylint: disable=protected-access
log.debug(f'Attention dispatcher: target={attn} previous={prev[0].value} active={current[0]} list={backends}')
elif len(attn) > 0:
log.warning(f'Attention dispatcher: active={prev[0].value} kernels={backends} target={attn} not found')
else:
log.debug(f'Attention dispatcher: active={prev[0].value} kernels={backends}')
def list_dispatcher_backends() -> list:
"""The kernels diffusers can dispatch attention to, for anything that offers hf_attention as a choice."""
try:
from diffusers.models import attention_dispatch as a
return sorted(b.value for b in a._AttentionBackendRegistry.list_backends()) # pylint: disable=protected-access
except Exception as e:
log.error(f'Attention dispatcher: {e}')
return []
modules/attention/registry.py
+108
View File
@@ -0,0 +1,108 @@
"""Declarative backend registry behind the scaled_dot_product_attention router."""
from dataclasses import dataclass, field
from typing import Callable
import torch
AttentionCall = Callable[..., torch.Tensor]
@dataclass(frozen=True)
class Platform:
"""Where the router runs: the devices backend name and the selected device."""
backend: str
device: torch.device | None = None
@dataclass(frozen=True)
class Constraints:
"""Shape, dtype and device conditions a backend serves; a call failing any of them moves on to the next entry."""
allow_cpu: bool = False
allow_mask: bool = True
allow_float32: bool = True
same_device: bool = False
head_dims: frozenset[int] | None = None
max_head_dim: int | None = None
min_tokens: int = 0 # query and key sequences both at least this long
min_long_side: int = 0 # query or key sequence longer than this
min_heads: int = 0
min_ndim: int = 0
def accepts(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, attn_mask: torch.Tensor | None) -> bool:
if not self.allow_cpu and query.device.type == 'cpu':
return False
if self.min_ndim and query.ndim < self.min_ndim:
return False
if not self.allow_mask and attn_mask is not None:
return False
if not self.allow_float32 and query.dtype == torch.float32:
return False
if self.same_device and (key.device != query.device or value.device != query.device):
return False
head_dim = query.shape[-1]
if self.head_dims is not None and head_dim not in self.head_dims:
return False
if self.max_head_dim is not None and head_dim > self.max_head_dim:
return False
if self.min_tokens and (query.shape[-2] < self.min_tokens or key.shape[-2] < self.min_tokens):
return False
if self.min_long_side and query.shape[-2] <= self.min_long_side and key.shape[-2] <= self.min_long_side:
return False
if self.min_heads and query.shape[-3] < self.min_heads:
return False
return True
@dataclass(frozen=True)
class AttentionBackend:
"""One attention implementation: how to prepare it once and which calls it serves."""
name: str
label: str # the cross_attention_optimization choice that enables it
priority: int # higher priority entries are tried first
prepare: Callable[[Platform, AttentionCall], AttentionCall | None] # imports and configures the implementation, returns its call or None
constraints: Constraints = field(default_factory=Constraints)
terminal: bool = False # serves every call the entries decline, in place of the original sdpa
platforms: frozenset[str] | None = None # devices backends the implementation exists for, None for all
options: tuple[str, ...] = () # settings the prepared call captures; a change to one rebuilds the chain
caps: frozenset[str] = frozenset() # what the call can consume beyond plain sdpa arguments: 'block_mask', and 'masked_block' when it composes one with a token mask
def available_on(self, platform: Platform) -> bool:
return self.platforms is None or platform.backend in self.platforms
def __repr__(self) -> str:
return f'AttentionBackend(name="{self.name}" label="{self.label}" priority={self.priority} terminal={self.terminal} platforms={list(self.platforms) if self.platforms is not None else []} options={self.options} caps={list(self.caps)})'
class Registry:
def __init__(self):
self.backends: dict[str, AttentionBackend] = {}
def register(self, backend: AttentionBackend) -> AttentionBackend:
if backend.name in self.backends:
raise ValueError(f'attention backend registered twice: name={backend.name}')
if self.by_label(backend.label) is not None:
raise ValueError(f'attention backend label registered twice: label="{backend.label}"')
self.backends[backend.name] = backend
return backend
def by_label(self, label: str) -> AttentionBackend | None:
return next((backend for backend in self.backends.values() if backend.label == label), None)
def ordered(self) -> list[AttentionBackend]:
"""Backends by ascending priority, the order they are prepared in."""
return sorted(self.backends.values(), key=lambda backend: backend.priority)
def labels(self) -> list[str]:
return [backend.label for backend in self.ordered()]
def options(self) -> list[str]:
return sorted({name for backend in self.backends.values() for name in backend.options})
def with_cap(self, cap: str) -> list[AttentionBackend]:
return [backend for backend in self.ordered() if cap in backend.caps]
def __repr__(self) -> str:
return f'Registry(backends={list(self.backends.keys())})'
registry = Registry()
modules/attention/router.py
+167
View File
@@ -0,0 +1,167 @@
"""The single scaled_dot_product_attention entry point over the prepared backends."""
from dataclasses import dataclass
from functools import wraps
from typing import Callable
import torch
from installer import torch_info
from modules.logger import log
from modules.attention import context, debug
from modules.attention.registry import AttentionBackend, AttentionCall, Platform, Registry, registry as default_registry
@dataclass(frozen=True)
class PlanEntry:
backend: AttentionBackend
call: AttentionCall
caps: frozenset[str] = frozenset() # the backend's declared caps, narrowed to what prepare verified in the installed implementation
@dataclass(frozen=True)
class Plan:
"""The prepared chain for one set of overrides: entries by descending priority, then the terminal or the original sdpa."""
entries: tuple[PlanEntry, ...]
terminal: PlanEntry | None
original: AttentionCall
platform: Platform
labels: tuple[str, ...]
def chain(self) -> list[str]:
names = [entry.backend.name for entry in self.entries]
names.append(self.terminal.backend.name if self.terminal is not None else 'sdpa')
return names
current_plan: Plan | None = None
def build_plan(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
reg = reg if reg is not None else default_registry
entries: list[PlanEntry] = []
terminal: PlanEntry | None = None
for backend in reg.ordered(): # ascending priority: the last prepared backend is tried first
if backend.label not in labels:
continue
if not backend.available_on(platform):
log.warning(f'Attention: type="{backend.label}" not available on backend={platform.backend}')
continue
try:
call = backend.prepare(platform, original)
except Exception as err:
log.error(f'Attention: type="{backend.label}" {err}')
continue
if call is None:
continue
entry = PlanEntry(backend=backend, call=call, caps=backend.caps & frozenset(getattr(call, 'caps', backend.caps)))
if backend.terminal:
terminal = entry
else:
entries.append(entry)
entries.reverse()
return Plan(entries=tuple(entries), terminal=terminal, original=original, platform=platform, labels=tuple(labels))
def make_router(plan: Plan, observer: Callable | None = None, stage: Callable | None = None) -> AttentionCall:
entries = plan.entries
terminal = plan.terminal.call if plan.terminal is not None else None
terminal_name = plan.terminal.backend.name if plan.terminal is not None else 'sdpa'
original = plan.original
@wraps(original)
def sdpa_router(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False, **kwargs):
for entry in entries:
if entry.backend.constraints.accepts(query, key, value, attn_mask):
if stage is not None and 'block_mask' in entry.caps:
selection = stage(query, key, value, attn_mask, is_causal, entry.caps)
if selection is not None:
if observer is not None:
observer(f'{entry.backend.name}+sparse', query, key, attn_mask)
return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa, selection=selection)
if observer is not None:
observer(entry.backend.name, query, key, attn_mask)
return entry.call(query, key, value, attn_mask, dropout_p, is_causal, scale, enable_gqa)
if observer is not None: # pylint: disable=duplicate-code
observer(terminal_name, query, key, attn_mask)
if terminal is not None:
return terminal(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, enable_gqa=enable_gqa, **kwargs)
if enable_gqa: # older sdpa signatures and platform wrappers reject the keyword, so it only travels when set
kwargs['enable_gqa'] = enable_gqa
return original(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs)
return sdpa_router
def install_router(labels, platform: Platform, original: AttentionCall, reg: Registry | None = None) -> Plan:
"""Prepare the enabled backends and install the router; an empty plan leaves the original sdpa in place."""
global current_plan # pylint: disable=global-statement
plan = build_plan(labels, platform, original, reg)
debug.reset()
observer = debug.observe if debug.enabled else None
stage = build_sparse_stage(plan)
torch.nn.functional.scaled_dot_product_attention = make_router(plan, observer, stage) if (plan.entries or plan.terminal is not None) else original
current_plan = plan
torch_info.set(attention='>'.join(plan.chain()))
log.debug(f'Attention: chain={">".join(plan.chain())} backend={platform.backend} sparse={stage is not None}')
return plan
def build_sparse_stage(plan: Plan):
"""Sparse attention is a stage over the chain rather than a chain member, so it needs a backend in the chain that consumes a block mask."""
from modules.attention.sparse import stage as sparse_stage
try:
options = sparse_stage.read_options()
except Exception:
return None
if not options.enabled:
return None
capable = [entry.backend.name for entry in plan.entries if 'block_mask' in entry.caps]
if not capable:
names = [backend.label for backend in default_registry.with_cap('block_mask')]
log.warning(f'Attention: sparse=True compatible={names} not set')
return None
built = sparse_stage.make_stage(options)
if built is not None:
log.info(f'Attention: sparse=True type={capable[0]} budget={options.budget:.0%} gate={options.min_tokens} schedule={options.schedule_steps}x+{options.schedule_bump:.0%}')
return built
def get_plan() -> Plan | None:
return current_plan
def reapply_options(reg: Registry | None = None) -> list[str]:
"""Settings whose change rebuilds the chain: the override set, the torch kernel flags, every option a backend captures, and the sparse stage."""
from modules.attention.sparse import stage as sparse_stage
reg = reg if reg is not None else default_registry
return ['sdp_options', 'cross_attention_optimization', *reg.options(), *sparse_stage.OPTION_NAMES]
def reapply() -> None:
"""Rebuild the chain from the current settings; a resident compiled model is reset so its graphs trace the new router."""
from modules import devices, shared
devices.set_sdpa_params()
compiled = getattr(shared, 'compiled_model_state', None)
if compiled is not None and getattr(compiled, 'is_compiled', False):
torch._dynamo.reset() # pylint: disable=protected-access
log.debug('Attention: dynamo=reset compiled model resident')
def report() -> dict:
"""The active chain, sparse stage and generation context, for the api and the debug log."""
from modules.attention.sparse import stage as sparse_stage
plan = current_plan
state = context.current
options = sparse_stage.read_options()
layout = state.layout
return {
'chain': plan.chain() if plan is not None else ['sdpa'],
'overrides': list(plan.labels) if plan is not None else [],
'sparse': {
'enabled': options.enabled,
'budget': options.budget,
'gate': options.min_tokens,
'capable': [entry.backend.name for entry in plan.entries if 'block_mask' in entry.caps] if plan is not None else [],
'layout': {'source': layout.source, 'kinds': list(layout.kinds()), 'length': layout.length} if layout is not None else None,
},
'backend': plan.platform.backend if plan is not None else None,
'context': {'active': state.active, 'role': state.role, 'step': state.step, 'steps': state.steps, 'model': state.model_key},
}
modules/attention/sparse/__init__.py
+8
View File
@@ -0,0 +1,8 @@
"""Block-sparse attention: the selector, the token layout it respects, and the consumers that apply it."""
from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks
from modules.attention.sparse.layout import Span, TokenLayout, block_pins, layout_from_index_kwargs, layout_from_kwargs, layout_from_prefix, layout_from_segments, publish_segments, segments_from_live
__all__ = [
'BlockSelection', 'SparseSpec', 'block_count', 'radial_blocks', 'schedule', 'select_blocks',
'Span', 'TokenLayout', 'block_pins', 'layout_from_index_kwargs', 'layout_from_kwargs', 'layout_from_prefix', 'layout_from_segments', 'publish_segments', 'segments_from_live',
]
modules/attention/sparse/flex.py
+39
View File
@@ -0,0 +1,39 @@
"""Turn a BlockSelection into the BlockMask FlexAttention consumes, and call it so the mask is honored."""
import torch
from torch.nn.attention.flex_attention import BlockMask, flex_attention, _dense_to_ordered
from modules.attention.sparse.selector import BlockSelection
compiled_flex = None
def to_block_mask(selection: BlockSelection, device=None) -> BlockMask:
"""All selected tiles go in the full slots, so mask_mod is never invoked and no dense S squared mask is built."""
keep = selection.keep
if device is not None and keep.device != device:
keep = keep.to(device)
if keep.dim() != 4:
raise ValueError(f'block selection must be 4d, got {tuple(keep.shape)}')
# the partial slots stay empty by construction, so build them directly rather than sorting a mask of zeros
empty_num = torch.zeros(keep.shape[:-1], dtype=torch.int32, device=keep.device)
empty_indices = torch.zeros(keep.shape, dtype=torch.int32, device=keep.device)
full_num, full_indices = _dense_to_ordered(keep)
return BlockMask.from_kv_blocks(
empty_num, empty_indices,
full_kv_num_blocks=full_num, full_kv_indices=full_indices,
BLOCK_SIZE=(selection.block_q, selection.block_kv),
seq_lengths=(selection.seq_q, selection.seq_kv), # exact lengths, so a ragged tail is handled rather than rounded up
compute_q_blocks=False, # backward only metadata, and inference never reads it
)
def flex_call():
"""flex_attention reads the block lists only when compiled; called eagerly it evaluates mask_mod instead and a block-only mask is silently dense."""
global compiled_flex # pylint: disable=global-statement
if compiled_flex is None:
compiled_flex = torch.compile(flex_attention, dynamic=False)
return compiled_flex
def attend(query, key, value, selection: BlockSelection, scale=None, enable_gqa=False):
return flex_call()(query, key, value, block_mask=to_block_mask(selection, device=query.device), scale=scale, enable_gqa=enable_gqa)
modules/attention/sparse/layout.py
+172
View File
@@ -0,0 +1,172 @@
"""What each token in a packed sequence is, so the selector knows what it may sparsify."""
from dataclasses import dataclass
import torch
# only the bulk modalities are sparsifiable; everything else is pinned dense, and an unrecognized kind pins too
SPARSIFIABLE = frozenset({'video', 'image'})
DROPPED = frozenset({'pad'})
@dataclass(frozen=True)
class Span:
kind: str
start: int
end: int
@dataclass(frozen=True)
class TokenLayout:
"""Ordered spans covering one packed sequence."""
spans: tuple[Span, ...]
length: int
source: str = 'unknown' # how the layout was obtained, for the log
def key(self) -> tuple:
return (self.length, self.source, tuple((s.kind, s.start, s.end) for s in self.spans))
def kinds(self) -> tuple[str, ...]:
return tuple(dict.fromkeys(s.kind for s in self.spans))
def sparsifiable_tokens(self) -> int:
return sum(s.end - s.start for s in self.spans if s.kind in SPARSIFIABLE)
def token_flags(self, device) -> tuple[torch.Tensor, torch.Tensor]:
"""Per token: may this be sparsified, and is it padding."""
sparse = torch.zeros(self.length, dtype=torch.bool, device=device)
pad = torch.zeros(self.length, dtype=torch.bool, device=device)
for span in self.spans:
if span.kind in SPARSIFIABLE:
sparse[span.start:span.end] = True
elif span.kind in DROPPED:
pad[span.start:span.end] = True
return sparse, pad
def runs(indices: torch.Tensor) -> list[tuple[int, int]]:
"""Contiguous [start, end) runs in a sorted 1d index tensor."""
if indices.numel() == 0:
return []
values = indices.detach().to('cpu', torch.int64).sort().values
breaks = (values[1:] - values[:-1] != 1).nonzero().flatten().tolist()
bounds = [0, *[b + 1 for b in breaks], values.numel()]
return [(int(values[bounds[i]].item()), int(values[bounds[i + 1] - 1].item()) + 1) for i in range(len(bounds) - 1)]
def layout_from_index_kwargs(kwargs: dict, length: int | None = None) -> TokenLayout | None:
"""Read a layout off the *_indices tensors a pipeline passes its transformer by name."""
spans: list[Span] = []
for name, value in kwargs.items():
if not name.endswith('_indices') or not torch.is_tensor(value) or value.dim() != 1 or value.is_floating_point():
continue
kind = name[:-len('_indices')].lower()
found = runs(value)
for position, (start, end) in enumerate(found):
# a video run that is not the last one is keyframe conditioning, which stays dense
resolved = 'cond' if (kind == 'video' and position < len(found) - 1) else kind
spans.append(Span(kind=resolved, start=start, end=end))
if not spans:
return None
spans.sort(key=lambda s: s.start)
return TokenLayout(spans=tuple(spans), length=length if length is not None else spans[-1].end, source='indices')
# how an architecture orders its joint sequence, which the call itself does not reveal. Verified against the
# diffusers transformers that take txt_ids and img_ids; HiDream packs image first and is deliberately absent, so
# it falls back rather than being pinned backwards. An unlisted class publishes nothing.
JOINT_TEXT_FIRST = frozenset({
'FluxTransformer2DModel', 'Flux2Transformer2DModel', 'ChromaTransformer2DModel', 'BriaTransformer2DModel',
'BriaFiboTransformer2DModel', 'LongCatImageTransformer2DModel', 'OvisImageTransformer2DModel',
})
def layout_from_stream_ids(kwargs: dict, cls_name: str | None) -> TokenLayout | None:
"""Read the stream lengths off the rotary id tensors a joint transformer is given by name."""
if cls_name not in JOINT_TEXT_FIRST:
return None
text, image = kwargs.get('txt_ids'), kwargs.get('img_ids')
if not torch.is_tensor(text) or not torch.is_tensor(image) or text.dim() < 2 or image.dim() < 2:
return None
return layout_from_segments((('text', text.shape[-2]), ('image', image.shape[-2])), source='stream-ids')
def layout_from_kwargs(kwargs: dict, cls_name: str | None = None) -> TokenLayout | None:
"""Whatever the denoiser says about its own packing, by whichever convention it uses."""
return layout_from_index_kwargs(kwargs or {}) or layout_from_stream_ids(kwargs or {}, cls_name)
def layout_from_segments(segments, length: int | None = None, source: str = 'segments') -> TokenLayout:
"""Build a layout from ordered (kind, count) pairs, the form a transformer knows at its packing site."""
spans: list[Span] = []
cursor = 0
for kind, count in segments:
if count <= 0:
continue
spans.append(Span(kind=kind, start=cursor, end=cursor + count))
cursor += count
return TokenLayout(spans=tuple(spans), length=length if length is not None else cursor, source=source)
def segments_from_live(live: torch.Tensor, kind: str, pad_kind: str = 'pad') -> list[tuple[str, int]]:
"""Run length encode a boolean live mask into ordered (kind, count) pairs, the dead runs labelled as padding."""
values = live.detach().to('cpu').bool()
if values.numel() == 0:
return []
changes = (values[1:] != values[:-1]).nonzero().flatten().tolist()
bounds = [0, *[c + 1 for c in changes], values.numel()]
return [(kind if bool(values[bounds[i]]) else pad_kind, bounds[i + 1] - bounds[i]) for i in range(len(bounds) - 1)]
def publish_segments(segments, length: int | None = None, source: str = 'segments') -> None:
"""Publish a layout from the site that packs the sequence, which is the only place the segment lengths are all known."""
from modules.attention import context
context.set_layout(layout_from_segments(segments, length=length, source=source))
def layout_from_prefix(length: int, prefix: int) -> TokenLayout:
"""Fallback when nothing published a layout: treat a leading run as conditioning and sparsify the rest."""
return layout_from_segments([('text', prefix), ('image', length - prefix)], length=length, source='prefix')
def block_flags(flags: torch.Tensor, block: int) -> tuple[torch.Tensor, torch.Tensor]:
"""Per block: do all tokens carry the flag, does any token carry it."""
seq = flags.shape[0]
whole = (seq // block) * block
parts_all, parts_any = [], []
if whole:
view = flags[:whole].view(whole // block, block)
parts_all.append(view.all(dim=-1))
parts_any.append(view.any(dim=-1))
if whole < seq:
parts_all.append(flags[whole:].all(dim=-1, keepdim=True))
parts_any.append(flags[whole:].any(dim=-1, keepdim=True))
def join(parts):
return parts[0] if len(parts) == 1 else torch.cat(parts, dim=0)
return join(parts_all), join(parts_any)
pin_cache: dict = {}
def block_pins(layout: TokenLayout, seq_q: int, seq_kv: int, block_q: int, block_kv: int, device) -> tuple[torch.Tensor, torch.Tensor]:
"""Tiles that must stay dense and tiles that can be skipped outright, as (1, 1, NQ, NK) masks."""
cache_key = (layout.key(), seq_q, seq_kv, block_q, block_kv, str(device))
hit = pin_cache.get(cache_key)
if hit is not None:
return hit
sparse_tokens, pad_tokens = layout.token_flags(device)
q_sparse = sparse_tokens[:seq_q] if layout.length >= seq_q else torch.nn.functional.pad(sparse_tokens, (0, seq_q - layout.length))
kv_sparse = sparse_tokens[:seq_kv] if layout.length >= seq_kv else torch.nn.functional.pad(sparse_tokens, (0, seq_kv - layout.length))
kv_pad = pad_tokens[:seq_kv] if layout.length >= seq_kv else torch.nn.functional.pad(pad_tokens, (0, seq_kv - layout.length))
q_all_sparse, _ = block_flags(q_sparse, block_q)
kv_all_sparse, _ = block_flags(kv_sparse, block_kv)
kv_all_pad, _ = block_flags(kv_pad, block_kv)
# a tile is pinned when its query tile or its key tile carries anything that is not sparsifiable, boundary tiles included
pins = (~q_all_sparse).unsqueeze(-1) | (~kv_all_sparse).unsqueeze(0)
drops = kv_all_pad.unsqueeze(0).expand_as(pins)
pins = (pins & ~drops).unsqueeze(0).unsqueeze(0).contiguous()
drops = drops.unsqueeze(0).unsqueeze(0).contiguous()
if len(pin_cache) > 32:
pin_cache.clear()
pin_cache[cache_key] = (pins, drops)
return pins, drops
modules/attention/sparse/selector.py
+148
View File
@@ -0,0 +1,148 @@
"""Fixed-budget block selection: which KV tiles each query tile attends to."""
from dataclasses import dataclass
import math
import torch
@dataclass(frozen=True)
class SparseSpec:
"""How much to keep and at what granularity. Budget is a fraction of the sparsifiable candidates, pins are added on top."""
budget: float = 0.30
block_q: int = 128
block_kv: int = 64
head_shared: bool = False # score once for all heads, cheaper and coarser
force: bool = False # skip the dense short circuit, so tests can exercise the path at budget 1.0
score_chunk_bytes: int = 256 << 20
@dataclass(frozen=True)
class BlockSelection:
"""int8 keep flags per (query tile, kv tile); the geometry every consumer reads."""
keep: torch.Tensor # (B, H, NQ, NK), H is the query head count or 1
block_q: int
block_kv: int
budget: float
seq_q: int
seq_kv: int
@property
def shape(self) -> tuple[int, int, int, int]:
b, h, nq, nk = self.keep.shape
return (b, h, nq, nk)
def density(self) -> float:
"""Fraction of tiles kept. Reads back from the accelerator, so this is for reporting and tests, never the hot path."""
return float(self.keep.sum().item()) / max(self.keep.numel(), 1)
def block_count(length: int, block: int) -> int:
return (length + block - 1) // block
def pool_blocks(x: torch.Tensor, block: int) -> torch.Tensor:
"""Mean over each block of tokens, fp32, without materializing a padded copy."""
seq = x.shape[-2]
whole = (seq // block) * block
parts = []
if whole:
head = x[..., :whole, :]
parts.append(head.unflatten(-2, (whole // block, block)).mean(dim=-2, dtype=torch.float32))
if whole < seq:
parts.append(x[..., whole:, :].mean(dim=-2, dtype=torch.float32, keepdim=True))
return parts[0] if len(parts) == 1 else torch.cat(parts, dim=-2)
def diagonal_blocks(nq: int, nk: int, block_q: int, block_kv: int, device) -> torch.Tensor:
"""Tiles whose query and key token ranges overlap; keeping them removes the empty-row case."""
q_index = torch.arange(nq, device=device).unsqueeze(-1)
k_index = torch.arange(nk, device=device).unsqueeze(0)
return (q_index * block_q < (k_index + 1) * block_kv) & (k_index * block_kv < (q_index + 1) * block_q)
def score_blocks(query: torch.Tensor, key: torch.Tensor, spec: SparseSpec) -> torch.Tensor:
"""Mean-pooled query-key affinity per tile pair. No scale and no softmax: top-k is invariant under both."""
pooled_q = pool_blocks(query, spec.block_q) # (B, Hq, NQ, D)
pooled_k = pool_blocks(key, spec.block_kv) # (B, Hkv, NK, D)
heads_q, heads_kv = pooled_q.shape[1], pooled_k.shape[1]
if spec.head_shared:
pooled_q = pooled_q.mean(dim=1, keepdim=True)
pooled_k = pooled_k.mean(dim=1, keepdim=True)
elif heads_kv != heads_q: # gqa: score on query heads, the geometry both consumers expect
pooled_k = pooled_k.repeat_interleave(heads_q // heads_kv, dim=1)
heads = pooled_q.shape[1]
per_head = pooled_q.shape[2] * pooled_k.shape[2] * 4
chunk = max(1, min(heads, spec.score_chunk_bytes // max(per_head, 1)))
if chunk >= heads:
return pooled_q @ pooled_k.transpose(-1, -2)
return torch.cat([pooled_q[:, i:i + chunk] @ pooled_k[:, i:i + chunk].transpose(-1, -2) for i in range(0, heads, chunk)], dim=1)
plan_cache: dict = {}
def selection_plan(spec: SparseSpec, nq: int, nk: int, pins, drops, device, cache_key=None):
"""The parts that depend only on geometry and layout, not on the tensors: what must be kept, what may be chosen, and how many."""
key = (cache_key, nq, nk, spec.block_q, spec.block_kv, spec.budget, str(device))
hit = plan_cache.get(key) if cache_key is not None else None
if hit is not None:
return hit
must = diagonal_blocks(nq, nk, spec.block_q, spec.block_kv, device).unsqueeze(0).unsqueeze(0)
if pins is not None:
must = must | pins
forbidden = drops if drops is not None else torch.zeros_like(must)
candidates = ~must & ~forbidden
per_row = candidates.sum(dim=-1, keepdim=True) # (.., NQ, 1)
keep_per_row = torch.ceil(per_row * spec.budget).to(torch.int64)
covers_everything = bool((keep_per_row >= per_row).all()) # one readback, amortized over the generation by the cache
built = (must, forbidden, candidates, keep_per_row, covers_everything)
if cache_key is not None:
if len(plan_cache) > 32:
plan_cache.clear()
plan_cache[key] = built
return built
def select_blocks(query: torch.Tensor, key: torch.Tensor, spec: SparseSpec, pins: torch.Tensor | None = None, drops: torch.Tensor | None = None, cache_key=None) -> BlockSelection | None:
"""Keep the highest scoring KV tiles per query tile within the budget, plus pins and the diagonal. None means attend densely."""
seq_q, seq_kv = query.shape[-2], key.shape[-2]
nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
device = query.device
must, forbidden, candidates, keep_per_row, covers_everything = selection_plan(spec, nq, nk, pins, drops, device, cache_key)
if covers_everything and not spec.force:
return None # the budget covers every candidate, so the mask would be dense
scores = score_blocks(query, key, spec)
scores = scores.masked_fill(~candidates.expand_as(scores), float('-inf'))
# rank rather than topk, so the per row budget varies without a host side k
order = scores.argsort(dim=-1, descending=True, stable=True)
rank = torch.empty_like(order)
rank.scatter_(-1, order, torch.arange(nk, device=device).expand_as(order))
keep = must | ((rank < keep_per_row) & candidates)
keep &= ~forbidden
return BlockSelection(keep=keep.to(torch.int8), block_q=spec.block_q, block_kv=spec.block_kv, budget=spec.budget, seq_q=seq_q, seq_kv=seq_kv)
def radial_blocks(seq_q: int, seq_kv: int, density: float, spec: SparseSpec, device) -> BlockSelection:
"""A band around the diagonal at the requested density: the static control the selector has to beat."""
nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
q_center = (torch.arange(nq, device=device).unsqueeze(-1) + 0.5) * spec.block_q
k_center = (torch.arange(nk, device=device).unsqueeze(0) + 0.5) * spec.block_kv
distance = (q_center - k_center).abs()
low, high = 0.0, float(max(seq_q, seq_kv))
for _ in range(40): # bisect the bandwidth, since the band width to density map has no closed form at the edges
mid = (low + high) / 2
if float((distance <= mid).to(torch.float32).mean().item()) < density:
low = mid
else:
high = mid
keep = (distance <= high).unsqueeze(0).unsqueeze(0).to(torch.int8)
return BlockSelection(keep=keep, block_q=spec.block_q, block_kv=spec.block_kv, budget=density, seq_q=seq_q, seq_kv=seq_kv)
def schedule(steps: int, budget: float, bump: float = 0.0, bump_steps: int = 0) -> tuple[float, ...]:
"""Per-step budgets, precomputed. At most two distinct values, so a compiled consumer sees at most two specializations."""
if bump <= 0 or bump_steps <= 0 or steps <= 0:
return tuple([budget] * max(steps, 0))
raised = min(1.0, budget + bump)
edge = min(bump_steps, math.ceil(steps / 2))
return tuple([raised if (i < edge or i >= steps - edge) else budget for i in range(steps)])
modules/attention/sparse/stage.py
+178
View File
@@ -0,0 +1,178 @@
"""The router stage that turns settings plus a published layout into a per call block selection."""
import os
from dataclasses import dataclass
import torch
from modules.logger import log
from modules.attention import context
from modules.attention.sparse import layout as layout_mod
from modules.attention.sparse.selector import BlockSelection, SparseSpec, block_count, radial_blocks, schedule, select_blocks
debug = os.environ.get('SD_ATTN_DEBUG', None) is not None
# SD_SPARSE_PATTERN=radial replaces the content aware selection with a static band around the
# diagonal at the same density: the control the selector has to beat, and the fallback if it does not
pattern = os.environ.get('SD_SPARSE_PATTERN', 'adaptive').strip().lower()
# measured on a 3090: below roughly this length a 30 percent budget caps under 1.25x per block,
# so the selector cannot pay for itself; see docs/sparse-attention-tracker.md. The settings registry
# carries the same number as the option default, this is the fallback when the option is absent
DEFAULT_MIN_TOKENS = 8192
# settings the stage reads, so a change to any of them rebuilds the chain
OPTION_NAMES = ('sparse_attention_enabled', 'sparse_attention_budget', 'sparse_attention_min_tokens', 'sparse_attention_schedule_steps', 'sparse_attention_schedule_bump', 'sparse_attention_head_shared', 'sparse_attention_exclude')
@dataclass(frozen=True)
class StageOptions:
enabled: bool = False
budget: float = 0.30
min_tokens: int = DEFAULT_MIN_TOKENS # 0 sparsifies every sequence that reaches the stage
schedule_steps: int = 0
schedule_bump: float = 0.0
head_shared: bool = False
exclude: tuple = () # architectures, pipeline classes or denoiser classes that stay dense
def read_options() -> StageOptions:
from modules import shared
opts = shared.opts
return StageOptions(
enabled=bool(getattr(opts, 'sparse_attention_enabled', False)),
budget=float(getattr(opts, 'sparse_attention_budget', 30)) / 100.0,
min_tokens=int(getattr(opts, 'sparse_attention_min_tokens', DEFAULT_MIN_TOKENS)),
schedule_steps=int(getattr(opts, 'sparse_attention_schedule_steps', 0)),
schedule_bump=float(getattr(opts, 'sparse_attention_schedule_bump', 0)) / 100.0,
head_shared=bool(getattr(opts, 'sparse_attention_head_shared', False)),
exclude=parse_exclusions(getattr(opts, 'sparse_attention_exclude', '')),
)
def parse_exclusions(raw) -> tuple:
"""The exclusion list as lowercase entries, each naming an architecture, a pipeline class or a denoiser class."""
return tuple(entry.strip().lower() for entry in str(raw or '').split(',') if len(entry.strip()) > 0)
def match_exclusion(model_key, arch, exclude: tuple) -> str:
"""The entry the loaded model matches, empty when none of them do."""
names = {str(name).lower() for name in (*(model_key or ()), arch) if name}
return next((entry for entry in exclude if entry in names), '')
def resolve_layout(seq: int, reported: set) -> layout_mod.TokenLayout:
"""The published layout when there is one, otherwise sparsify the whole sequence and say so once."""
published = context.current.layout
if isinstance(published, layout_mod.TokenLayout) and published.length == seq:
return published
if seq not in reported:
reported.add(seq)
detail = 'none published' if published is None else f'published length {getattr(published, "length", None)} does not match {seq}'
log.info(f'Sparse attention: no token layout ({detail}), sparsifying the whole sequence at tokens={seq}')
return layout_mod.layout_from_prefix(seq, 0)
def make_stage(options: StageOptions):
"""Return the per call selector, or None when the feature is off."""
if not options.enabled or options.budget >= 1.0:
return None
reported: set = set()
inactive: set = set()
notified: set = set()
cache: dict = {}
static: dict = {}
excluded: dict = {}
def static_selection(query, key, spec, pins, drops, cache_key):
"""A density matched band, built once per geometry, honoring the same layout pins so the control differs from the selector only in how it chooses video tiles."""
static_key = (cache_key, spec.budget, query.shape[-2], key.shape[-2])
built = static.get(static_key)
if built is None:
reference = select_blocks(query, key, spec, pins=pins, drops=drops, cache_key=cache_key)
if reference is None:
return None
target = reference.density()
pinned = float(pins.to(torch.float32).mean().item()) if pins is not None else 0.0
band = radial_blocks(query.shape[-2], key.shape[-2], max(target - pinned, 0.0), spec, query.device)
keep = band.keep.bool()
if pins is not None:
keep = keep | pins
if drops is not None:
keep = keep & ~drops
built = BlockSelection(keep=keep.to(torch.int8), block_q=spec.block_q, block_kv=spec.block_kv, budget=spec.budget, seq_q=query.shape[-2], seq_kv=key.shape[-2])
static.clear()
static[static_key] = built
log.info(f'Sparse attention: static radial pattern density={built.density():.3f} against selector {target:.3f} at budget={spec.budget:.0%}')
return built
def budget_for_step() -> float:
state = context.current
if options.schedule_steps <= 0 or options.schedule_bump <= 0 or state.steps <= 0:
return options.budget
key = (state.steps, options.budget, options.schedule_bump, options.schedule_steps)
table = cache.get(key)
if table is None:
table = schedule(state.steps, options.budget, options.schedule_bump, options.schedule_steps)
cache.clear()
cache[key] = table
return table[min(state.step, len(table) - 1)] if table else options.budget
def on_exclusion_list(state) -> bool:
"""Whether the loaded model is excluded, resolved once per model since the answer cannot change within one."""
if not options.exclude:
return False
hit = excluded.get(state.model_key)
if hit is None:
from modules import shared
hit = match_exclusion(state.model_key, getattr(shared, 'sd_model_type', None), options.exclude)
excluded[state.model_key] = hit
if hit and debug: # an enabled setting that cannot act says so rather than doing nothing quietly
log.trace(f'Sparse attention: "{hit}" is on the exclusion list; attention stays dense')
return len(hit) > 0
def decline(reason: str):
stage.last_skip = reason
return None
def stage(query, key, value, attn_mask, is_causal, caps=frozenset()): # pylint: disable=unused-argument
state = context.current
if state.role != 'transformer' or not state.active:
return decline('not the denoiser')
if on_exclusion_list(state):
return decline('excluded')
if is_causal: # the selection keeps the diagonal but encodes no causality
return decline('causal')
if attn_mask is not None and 'masked_block' not in caps: # flex would need a mask_mod to combine the two
if 'masked' not in notified: # an enabled setting that cannot act says so rather than doing nothing quietly
notified.add('masked')
if debug:
log.trace('Sparse attention: this model passes an attention mask and the serving backend cannot combine it with a block selection; attention stays dense')
return decline('masked')
if query.device.type == 'cpu' or query.dim() != 4:
return decline('unsupported tensor')
seq_q, seq_kv = query.shape[-2], key.shape[-2]
if seq_q != seq_kv: # cross attention is short and already cheap
return decline('cross attention')
if seq_q < options.min_tokens:
if seq_q not in inactive: # an enabled setting that cannot act says so rather than doing nothing quietly
inactive.add(seq_q)
if debug:
log.trace(f'Sparse attention: inactive at tokens={seq_q}, below the minimum sequence of {options.min_tokens}; attention stays dense')
return decline('below the minimum sequence')
budget = budget_for_step()
if budget >= 1.0:
return decline('budget covers everything')
spec = SparseSpec(budget=budget, head_shared=options.head_shared)
token_layout = resolve_layout(seq_q, reported)
nq, nk = block_count(seq_q, spec.block_q), block_count(seq_kv, spec.block_kv)
pins, drops = layout_mod.block_pins(token_layout, seq_q, seq_kv, spec.block_q, spec.block_kv, query.device)
if pins.shape[-2:] != (nq, nk):
return decline('layout geometry mismatch')
stage.last_skip = None
if pattern == 'radial':
return static_selection(query, key, spec, pins, drops, token_layout.key())
return select_blocks(query, key, spec, pins=pins, drops=drops, cache_key=token_layout.key())
stage.options = options
stage.last_skip = None
return stage
modules/ben2/__init__.py
-24
View File
@@ -1,24 +0,0 @@
model = None
def remove(image, refine: bool = True):
global model # pylint: disable=global-statement
from modules import shared, devices
if model is None:
from huggingface_hub import hf_hub_download
from .ben2_model import BEN_Base
model = BEN_Base()
model_file = hf_hub_download(
repo_id='PramaLLC/BEN2',
filename='BEN2_Base.pth',
cache_dir=shared.opts.hfcache_dir)
model.loadcheckpoints(model_file)
model = model.to(device=devices.device, dtype=devices.dtype).eval()
model = model.to(device=devices.device)
foreground = model.inference(image, refine_foreground=refine)
model = model.to(device=devices.cpu)
if foreground is None:
return image
return foreground
modules/ben2/ben2_model.py
-1298
View File
File diff suppressed because it is too large Load Diff
modules/caption/gemini.py
+11 -8
View File
@@ -3,6 +3,7 @@ import os
from modules import shared
from modules.logger import log
debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
@@ -11,7 +12,7 @@ class GoogleGeminiPipeline():
def __init__(self, model_name: str):
self.model = model_name.split(' (')[0]
from installer import install
install('google-genai==1.52.0')
install('google-genai==2.22.0')
from google import genai # pylint: disable=no-name-in-module
args = self.get_args()
self.client = genai.Client(**args)
@@ -63,11 +64,13 @@ class GoogleGeminiPipeline():
config['temperature'] = kwargs['temperature']
if 'max_output_tokens' in kwargs:
config['max_output_tokens'] = kwargs['max_output_tokens']
debug_log(f'Gemini config: {config}')
debug_log(f'LLM config: {config}')
debug_log(f'LLM instructions: "{instructions}"')
debug_log(f'LLM image: {image}')
question = question.replace('<', '').replace('>', '').replace('_', ' ')
if prefill:
question += prefill
debug_log(f'Gemini question: "{question}"')
debug_log(f'LLM question: "{question}"')
if image:
data = io.BytesIO()
@@ -80,9 +83,9 @@ class GoogleGeminiPipeline():
answer = ''
try:
response = self.client.models.generate_content(
model=model,
contents=contents,
config=config,
model = model,
contents = contents,
config = config,
)
debug_log(f'Gemini response: {response}')
answer = response.text
@@ -94,8 +97,8 @@ class GoogleGeminiPipeline():
ai = None
def predict(question, image, vqa_model, system_prompt, model_name, prefill, thinking, gen_kwargs):
def predict(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs):
global ai # pylint: disable=global-statement
if ai is None:
ai = GoogleGeminiPipeline(model_name)
return ai(question, image, vqa_model, system_prompt, prefill, thinking, gen_kwargs)
return ai(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs)
modules/caption/grok.py
+99
View File
@@ -0,0 +1,99 @@
import io
import os
import base64
from modules import shared
from modules.logger import log
debug_enabled = os.environ.get('SD_CAPTION_DEBUG', None) is not None
debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
class XAIGrokPipeline():
def __init__(self, model_name: str):
self.url = 'https://api.x.ai/v1'
self.model = model_name.split(' (')[0].replace('xai/', '')
from installer import install
install('openai')
from openai import OpenAI # pylint: disable=no-name-in-module
args = self.get_args()
if not args:
return
self.client = OpenAI(**args)
log.debug(f'Load model: type=XAIGrok model="{self.model}"')
def get_args(self):
from modules.shared import opts
# Use UI settings only - env vars are intentionally ignored
api_key = opts.xai_api_key
has_api_key = api_key and len(api_key) > 0
if not has_api_key: # Gemini Developer API: api_key only
log.error(f'Cloud: model="{self.model}" API key not provided')
return None
args = {
'api_key': api_key,
'base_url': self.url,
}
# Debug logging
args_log = args.copy()
if args_log.get('api_key'):
args_log['api_key'] = '...' + args_log['api_key'][-4:]
log.debug(f'Cloud: model="{self.model}" args={args_log}')
return args
def __call__(self, question, image, model, instructions, prefill, thinking, kwargs):
question = question.replace('<', '').replace('>', '').replace('_', ' ')
if prefill:
question += prefill
debug_log(f'LLM instructions: "{instructions}"')
debug_log(f'LLM question: "{question}"')
debug_log(f'LLM image: {image}')
answer = ''
temperature = kwargs.get('temperature', 0.0)
try:
if image is not None:
image_data = image.convert('RGB')
image_bytes = io.BytesIO()
image_data.save(image_bytes, format='JPEG')
image_bytes.seek(0)
content = [
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{base64.b64encode(image_bytes.getvalue()).decode('utf-8')}",
"detail": "high",
},
},
{
"type": "text",
"text": question,
},
]
else:
content = question
response = self.client.chat.completions.create(
model = self.model,
messages = [
{"role": "system", "content": instructions or shared.opts.caption_vlm_system},
{"role": "user", "content": content},
],
stream = False,
temperature = temperature,
reasoning_effort = "high" if thinking else "low"
)
text = (response.choices[0].message.content or "").strip()
debug_log(f'Grok response: {response}')
answer = text
except Exception as e:
log.error(f'Grok: {e}')
answer = f'Error: {e}'
return answer
ai = None
def predict(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs):
global ai # pylint: disable=global-statement
if ai is None:
ai = XAIGrokPipeline(model_name)
return ai(question, image, model_name, system_prompt, prefill, thinking, gen_kwargs)
modules/caption/models_def.py
+16 -7
View File
@@ -74,13 +74,17 @@ vlm_models = {
"AIDC Ovis2 2B": "AIDC-AI/Ovis2-2B",
"AIDC Ovis2 1B": "AIDC-AI/Ovis2-1B",
# cloud
f"Google Gemini 3.5 Flash {ui_symbols.cloud}": "google/gemini-3.5-flash",
f"Google Gemini 3.1 Pro {ui_symbols.cloud}": "gemini-3.1-pro-preview",
f"Google Gemini 3.8 Flash {ui_symbols.cloud}": "gemini-3.8-flash",
f"Google Gemini 3.7 Flash {ui_symbols.cloud}": "gemini-3.7-flash",
f"Google Gemini 3.6 Flash {ui_symbols.cloud}": "gemini-3.6-flash",
f"Google Gemini 3.5 Flash {ui_symbols.cloud}": "gemini-3.5-flash",
f"Google Gemini 3.5 Flash Lite {ui_symbols.cloud}": "gemini-3.5-flash-lite",
f"Google Gemini 3.1 Flash Lite {ui_symbols.cloud}": "gemini-3.1-flash-lite",
f"Google Gemini 3.1 Flash Lite Preview {ui_symbols.cloud}": "gemini-3.1-flash-lite-preview",
f"Google Gemini 2.5 Pro {ui_symbols.cloud}": "gemini-2.5-pro",
f"Google Gemini 2.5 Flash {ui_symbols.cloud}": "gemini-2.5-flash",
f"Google Gemini 2.5 Flash Lite {ui_symbols.cloud}": "gemini-2.5-flash-lite",
f"Google Gemini 3.1 Pro {ui_symbols.cloud}": "gemini-3.1-pro-preview",
f"X.AI Grok 3 {ui_symbols.cloud}": "grok-3-latest",
f"X.AI Grok 3 Fast {ui_symbols.cloud}": "grok-3-fast-latest",
f"X.AI Grok 3 Mini {ui_symbols.cloud}": "grok-3-mini-latest",
f"X.AI Grok 3 Mini Fast {ui_symbols.cloud}": "grok-3-mini-fast-latest",
}
# Default model
@@ -224,5 +228,10 @@ Summary:
def get_vlm_repo(display_name: str) -> str:
"""Look up repo ID from display name, stripping any trailing symbols."""
from modules.logger import log
name = display_name.strip()
return vlm_models.get(name, name)
model = vlm_models.get(name, None)
if model is None:
log.warning(f"Model '{name}' not found")
return name
return model
modules/caption/vqa.py
+6 -1
View File
@@ -1597,7 +1597,12 @@ class VQA:
handler = 'gemini'
gen_kwargs = get_kwargs(self.model)
from modules.caption import gemini
answer = gemini.predict(question, image, vqa_model, system_prompt, model_name, prefill, thinking_mode, gen_kwargs)
answer = gemini.predict(question, image, vqa_model, system_prompt, prefill, thinking_mode, gen_kwargs)
elif 'grok' in vqa_model.lower():
handler = 'grok'
gen_kwargs = get_kwargs(self.model)
from modules.caption import grok
answer = grok.predict(question, image, vqa_model, system_prompt, prefill, thinking_mode, gen_kwargs)
else:
answer = 'unknown model'
except Exception as e:
modules/caption/vqa_detection.py
+1 -1
View File
@@ -313,7 +313,7 @@ def draw_bounding_boxes(image: Image.Image, detections: list, points: list | Non
# Try to load a font, fall back to default if unavailable
try:
font_size = max(12, int(min(width, height) * 0.02))
font_path = shared.opts.font or os.path.join(script_path, "ui", "fonts", "notosans-nerdfont-regular.ttf")
font_path = shared.opts.font or os.path.join(script_path, "ui", "css", "ubuntu-nerdfont.ttf")
font = ImageFont.truetype(font_path, size=font_size)
except Exception:
font = ImageFont.load_default()
modules/civitai/api_civitai.py
+76 -50
View File
@@ -71,6 +71,8 @@ def file_to_legacy_dict(f) -> dict:
'size': int(f.size_kb * 1024),
'name': f.name,
'type': f.type,
'primary': bool(f.primary),
'metadata': {'fp': f.metadata.fp, 'format': f.metadata.format, 'size': f.metadata.size, 'quantType': f.metadata.quant_type},
'hashes': [h for h in [f.hashes.sha256, f.hashes.autov1, f.hashes.autov2, f.hashes.autov3, f.hashes.crc32, f.hashes.blake3] if h],
'url': f.download_url,
}
@@ -158,6 +160,15 @@ def get_version_by_hash(hash_str: str, token: str | None = None):
return version_to_dict(version)
def get_version_mini(version_id: int, token: str | None = None):
"""Download-shaped version view carrying permission and early-access flags."""
from modules.civitai.client_civitai import client
version = client.get_version_mini(version_id, token=token)
if version is None:
return JSONResponse(content={"error": "version not found"}, status_code=404)
return version_to_dict(version)
def get_options():
"""Get valid types, sort, period, base_models from CivitAI API discovery."""
from modules.civitai.client_civitai import client
@@ -374,14 +385,14 @@ def post_settings(request: dict):
if user is None:
return JSONResponse(content={"error": "Invalid API token"}, status_code=400)
log.info(f'CivitAI token validated: user={user.get("username", "?")}')
shared.opts.data['civitai_token'] = token.strip()
shared.opts.civitai_token = token.strip()
save_subfolder_enabled = request.get('save_subfolder_enabled')
if save_subfolder_enabled is not None:
shared.opts.data['civitai_save_subfolder_enabled'] = bool(save_subfolder_enabled)
shared.opts.civitai_save_subfolder_enabled = bool(save_subfolder_enabled)
if save_subfolder is not None:
shared.opts.data['civitai_save_subfolder'] = save_subfolder
shared.opts.civitai_save_subfolder = save_subfolder
if discard_hash_mismatch is not None:
shared.opts.data['civitai_discard_hash_mismatch'] = discard_hash_mismatch
shared.opts.civitai_discard_hash_mismatch = discard_hash_mismatch
shared.opts.save()
return get_settings()
@@ -415,9 +426,12 @@ def post_metadata_scan(request: dict | None = None):
from modules.civitai import metadata_civitai
page = (request or {}).get('page', None)
results = []
for batch in metadata_civitai.civit_search_metadata(title=page, raw=True):
if isinstance(batch, list):
results = batch
try:
for batch in metadata_civitai.civit_search_metadata(title=page, raw=True):
if isinstance(batch, list):
results = batch
except metadata_civitai.SweepBusy as e:
return JSONResponse(content={"error": str(e)}, status_code=409)
return {"results": results}
@@ -425,9 +439,12 @@ def post_metadata_update():
"""Update local metadata from CivitAI."""
from modules.civitai import metadata_civitai
items = []
for batch in metadata_civitai.civit_update_metadata(raw=True):
if isinstance(batch, list):
items = batch
try:
for batch in metadata_civitai.civit_update_metadata(raw=True):
if isinstance(batch, list):
items = batch
except metadata_civitai.SweepBusy as e:
return JSONResponse(content={"error": str(e)}, status_code=409)
results = []
for item in items:
results.append({
@@ -550,7 +567,6 @@ def buildsidecar_index():
continue
# Match the companion file to a JSON entry by size (sizeKB)
companion_size_kb = os.path.getsize(companion) / 1024.0
companion_name = os.path.basename(base)
best_sha = None
best_diff = float('inf')
for v in data.get('modelVersions', []):
@@ -563,7 +579,7 @@ def buildsidecar_index():
best_sha = sha
best_diff = diff
if best_sha:
sidecar_index[best_sha.lower()] = {"filename": companion_name, "type": model_type}
sidecar_index[best_sha.lower()] = {"filename": companion, "type": model_type}
except Exception:
continue
log.debug(f'CivitAI sidecar index: {len(sidecar_index)} hashes from sidecar files')
@@ -580,54 +596,60 @@ def invalidatesidecar_index():
# ---------------------------------------------------------------------------
def post_check_local(request: dict):
"""Check which SHA256 hashes correspond to locally downloaded files."""
"""Check which SHA256 hashes correspond to local model files, dropping hash cache entries whose files are gone."""
from modules import hashes as hash_module
input_hashes = request.get('hashes', [])
if not input_hashes:
from modules.civitai.filemanage_civitai import hash_cache_path, prune_hash_cache
requested = [str(h) for h in request.get('hashes', []) if h]
if not requested:
return {"found": {}}
# Build reverse lookup: lowercase sha256 -> {filename, type}
prune_hash_cache()
wanted = {h.lower() for h in requested}
titles_by_sha: dict[str, list[str]] = {}
for title, entry in list(hash_module.cache().items()):
sha = (entry.get("sha256") or "").lower()
if sha in wanted:
titles_by_sha.setdefault(sha, []).append(title)
found = {}
for title, entry in hash_module.cache().items():
sha = entry["sha256"]
if not sha:
continue
parts = title.split("/", 1)
file_type = parts[0] if len(parts) > 1 else "unknown"
found[sha.lower()] = {"filename": title, "type": file_type}
# Supplement from in-memory checkpoint registry
gone = []
for sha, titles in titles_by_sha.items():
for title in titles:
path = hash_cache_path(title)
if path is None: # no loaded registry names the file
continue
if os.path.exists(path):
found[sha] = {"filename": path, "type": title.split("/", 1)[0]}
break
gone.append(title)
if gone:
for title in gone:
hash_module.cache().pop(title, None)
hash_module.save_cache()
log.debug(f'CivitAI check local: pruned={len(gone)} hash cache entries without files')
try:
from modules.sd_checkpoint import checkpoints_list
for _title, cp in checkpoints_list.items():
if cp.sha256:
key = cp.sha256.lower()
if key not in found:
found[key] = {"filename": cp.filename, "type": "checkpoint"}
for cp in checkpoints_list.values():
key = (cp.sha256 or "").lower()
if key in wanted and key not in found and os.path.exists(cp.filename):
found[key] = {"filename": cp.filename, "type": "checkpoint"}
except Exception:
pass
# Supplement from in-memory LoRA registry
try:
from modules.lora.lora_load import available_networks
for _name, net in available_networks.items():
if net.hash:
key = net.hash.lower()
if key not in found:
found[key] = {"filename": net.filename, "type": "lora"}
for net in available_networks.values():
key = (net.hash or "").lower()
if key in wanted and key not in found and os.path.isfile(net.filename):
found[key] = {"filename": net.filename, "type": "lora"}
except Exception:
pass
# Supplement from sidecar index (covers files never hashed locally)
sidecar = buildsidecar_index()
for h in input_hashes:
if not h:
continue
key = h.lower()
if key not in found and key in sidecar:
found[key] = sidecar[key]
# Match requested hashes
result = {}
for h in input_hashes:
if not h:
continue
match = found.get(h.lower())
for h in requested:
key = h.lower()
match = found.get(key)
if match is None:
entry = sidecar.get(key)
if entry and os.path.isfile(entry["filename"]):
match = entry
if match:
result[h] = match
return {"found": result}
@@ -664,7 +686,7 @@ def legacy_get_civitai(
query=query, tag=tag, types=types, sort=sort, period=period,
nsfw=nsfw, limit=limit, base=base, token=token, exact=exact,
)
return [model_to_legacy_dict(m) for m in models]
return [model_to_legacy_dict(m) for m in models.items]
return JSONResponse(content=[], status_code=200)
@@ -672,8 +694,11 @@ def legacy_post_civitai(page: str | None = None):
"""Legacy POST /sdapi/v1/civitai — scan metadata."""
from modules.civitai import metadata_civitai
result = []
for r in metadata_civitai.civit_search_metadata(title=page, raw=True):
result = r
try:
for r in metadata_civitai.civit_search_metadata(title=page, raw=True):
result = r
except metadata_civitai.SweepBusy as e:
return JSONResponse(content={"error": str(e)}, status_code=409)
return result
@@ -687,6 +712,7 @@ def register_api(api):
api.add_api_route("/sdapi/v2/civitai/model/{model_id}", get_model, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/version/{version_id}", get_version, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/version/by-hash/{hash_str}", get_version_by_hash, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/version/mini/{version_id}", get_version_mini, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/options", get_options, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/tags", get_tags, methods=["GET"], tags=["CivitAI"])
api.add_api_route("/sdapi/v2/civitai/creators", get_creators, methods=["GET"], tags=["CivitAI"])
modules/civitai/basemodels_civitai.py
-134
View File
@@ -1,134 +0,0 @@
import re
import time
from modules.logger import log
# Canonical base-model metadata lives in the civitai/civitai repo. The live
# /images validator is generated from it, so the file has metadata (group,
# ecosystem, engine, hidden) while the validator has the current name list.
# Callers merge both.
github_cache: list[dict] = []
github_cache_time: float = 0
GITHUB_TTL = 6 * 3600 # 6 hours
GITHUB_URL = 'https://raw.githubusercontent.com/civitai/civitai/main/src/shared/constants/base-model.constants.ts'
def parse_base_model_config(ts_source: str) -> list[dict]:
"""Parse the baseModelConfig array from base-model.constants.ts.
Uses character-by-character bracket walking rather than regex because
entries can span multiple lines. Returns list of dicts with keys
name/type/group/hidden plus optional ecosystem/engine/family.
"""
start_match = re.search(r'const\s+baseModelConfig\s*=\s*\[', ts_source)
if not start_match:
return []
# Walk to the matching ] respecting string literals
pos = start_match.end()
depth = 1
in_string: str | None = None
end_pos = -1
while pos < len(ts_source):
ch = ts_source[pos]
if in_string is not None:
if ch == '\\':
pos += 2
continue
if ch == in_string:
in_string = None
else:
if ch in ("'", '"', '`'):
in_string = ch
elif ch == '[':
depth += 1
elif ch == ']':
depth -= 1
if depth == 0:
end_pos = pos
break
pos += 1
if end_pos < 0:
return []
array_body = ts_source[start_match.end():end_pos]
# Extract top-level {...} entries, respecting strings and nested braces
entries: list[str] = []
brace_start = -1
brace_depth = 0
in_string = None
i = 0
while i < len(array_body):
ch = array_body[i]
if in_string is not None:
if ch == '\\':
i += 2
continue
if ch == in_string:
in_string = None
else:
if ch in ("'", '"', '`'):
in_string = ch
elif ch == '{':
if brace_depth == 0:
brace_start = i
brace_depth += 1
elif ch == '}':
brace_depth -= 1
if brace_depth == 0 and brace_start >= 0:
entries.append(array_body[brace_start:i + 1])
brace_start = -1
i += 1
# Per-entry field extraction (string + bool values only)
field_re = re.compile(
r"(\w+)\s*:\s*(?:'([^'\\]*(?:\\.[^'\\]*)*)'|\"([^\"\\]*(?:\\.[^\"\\]*)*)\"|(true|false))"
)
parsed: list[dict] = []
for entry in entries:
fields: dict = {}
for m in field_re.finditer(entry):
key = m.group(1)
if m.group(2) is not None:
fields[key] = m.group(2)
elif m.group(3) is not None:
fields[key] = m.group(3)
elif m.group(4) is not None:
fields[key] = m.group(4) == 'true'
if 'name' in fields and 'type' in fields and 'group' in fields:
item: dict = {
'name': fields['name'],
'type': fields['type'],
'group': fields['group'],
'hidden': bool(fields.get('hidden', False)),
}
for opt in ('ecosystem', 'engine', 'family'):
if opt in fields:
item[opt] = fields[opt]
parsed.append(item)
return parsed
def fetch_github_base_models() -> list[dict]:
"""Fetch and parse civitai's base-model constants from GitHub.
Returns list of metadata dicts (name, type, group, hidden, plus
optional ecosystem/engine/family). Returns [] on any failure;
callers fall back to the live /images probe. Cached with a longer
TTL than discover_options since these constants change rarely.
"""
global github_cache, github_cache_time # pylint: disable=global-statement
now = time.time()
if github_cache and (now - github_cache_time) < GITHUB_TTL:
return github_cache
try:
from modules import shared
r = shared.req(GITHUB_URL)
if r.status_code != 200:
log.debug(f'CivitAI github constants: code={r.status_code}')
return []
parsed = parse_base_model_config(r.text)
if parsed:
github_cache = parsed
github_cache_time = now
return parsed
except Exception as e:
log.debug(f'CivitAI github constants fetch failed: {e}')
return []
modules/civitai/client_civitai.py
+164 -21
View File
@@ -1,8 +1,10 @@
import os
import json
import time
import threading
from types import SimpleNamespace
from modules.logger import log
from modules.civitai.basemodels_civitai import fetch_github_base_models
from modules.civitai.models_civitai import CivitModel, CivitVersion, CivitImage, CivitSearchResponse, CivitTagResponse, CivitCreatorResponse, CivitUserProfile
from modules.civitai.models_civitai import CivitModel, CivitVersion, CivitVersionMini, CivitImage, CivitSearchResponse, CivitTagResponse, CivitCreatorResponse, CivitUserProfile
options_cache: dict = {}
@@ -11,6 +13,64 @@ OPTIONS_TTL = 3600 # 1 hour
# Civitai nsfwLevel bitmask: 1=PG/None 2=PG-13/Soft 4=R/Mature 8=X 16=XXX 32=Blocked
NSFW_LEVEL_SFW = 3 # None + Soft: Civitai's SFW browsing boundary
NSFW_LEVEL_ALL = 63 # every level set: disables filtering
BY_HASH_IDS_LIMIT = 10000 # POST /model-versions/by-hash/ids request cap
BY_HASH_LIMIT = 100 # POST /model-versions/by-hash request cap
MODEL_IDS_LIMIT = 100 # GET /models page cap; longer ids lists paginate
RETRY_LIMIT = 4 # retries after HTTP 429
RETRY_DELAY_MAX = 60 # seconds
request_slots: threading.BoundedSemaphore | None = None
request_slots_lock = threading.Lock()
def get_request_slots() -> threading.BoundedSemaphore:
"""Process-wide cap on concurrent CivitAI API requests, sized to shared.max_workers."""
global request_slots # pylint: disable=global-statement
with request_slots_lock:
if request_slots is None:
from modules.shared import max_workers
request_slots = threading.BoundedSemaphore(max_workers)
return request_slots
def retry_delay(response, attempt: int) -> float:
"""Seconds before retrying a 429: Retry-After when given in seconds, otherwise exponential."""
headers = getattr(response, 'headers', None) or {}
try:
delay = float(headers.get('Retry-After'))
except (TypeError, ValueError):
delay = 2 ** attempt
return min(max(delay, 0.0), RETRY_DELAY_MAX)
def response_message(response) -> str:
"""CivitAI error text from a failed response: its error string, ZodError issues, or the HTTP reason."""
try:
body = response.json()
except Exception:
body = None
error = body.get('error') if isinstance(body, dict) else None
if isinstance(error, dict):
error = error.get('message', '')
try:
error = '; '.join(f"{'.'.join(str(p) for p in issue.get('path', []))}: {issue.get('message', '')}" for issue in json.loads(error))
except Exception:
pass
message = body.get('message') if isinstance(body, dict) else None
if isinstance(error, str) and isinstance(message, str) and message and message != error: # download refusals carry a short error and a longer message
error = f'{error}: {message}'
if not error:
error = getattr(response, 'reason', '') or getattr(response, 'text', '')
return str(error).strip()[:200]
def post_json(url: str, body, headers: dict):
"""POST with the timeout, TLS and failure shape of shared.req."""
import requests
try:
return requests.post(url, json=body, timeout=30, headers=headers, verify=False, allow_redirects=True)
except Exception as e:
log.error(f'HTTP request error: url={url} {e}')
return SimpleNamespace(status_code=500, text=f'HTTP request error: url={url} {e}')
class CivitaiClient:
@@ -25,7 +85,7 @@ class CivitaiClient:
return tok
return os.environ.get('CIVITAI_TOKEN', None)
def _get(self, path: str, params: dict | None = None, token: str | None = None, stream: bool = False):
def send(self, method: str, path: str, params: dict | None = None, body=None, token: str | None = None, stream: bool = False):
from modules import shared
url = f"{self.BASE_URL}{path}"
headers = {}
@@ -37,7 +97,23 @@ class CivitaiClient:
query = urlencode({k: v for k, v in params.items() if v is not None and v != ''}, doseq=True)
if query:
url = f"{url}?{query}"
return shared.req(url, headers=headers if headers else None, stream=stream)
attempt = 0
while True:
with get_request_slots():
if method == 'POST':
r = post_json(url, body, headers)
else:
r = shared.req(url, headers=headers if headers else None, stream=stream)
retry_after = (getattr(r, 'headers', None) or {}).get('Retry-After')
if not (r.status_code == 429 or (r.status_code == 503 and retry_after is not None)) or attempt >= RETRY_LIMIT: # CivitAI sends 503 with Retry-After when search is overloaded
return r
delay = retry_delay(r, attempt)
log.warning(f'CivitAI retry: path={path} code={r.status_code} attempt={attempt + 1} delay={delay:.0f}s message="{response_message(r)}"')
time.sleep(delay)
attempt += 1
def _get(self, path: str, params: dict | None = None, token: str | None = None, stream: bool = False):
return self.send('GET', path, params=params, token=token, stream=stream)
def search_models(self, *, query: str = "", tag: str = "", types: str = "", sort: str = "", period: str = "",
base_models: list[str] | None = None, nsfw: bool | None = None, limit: int = 20,
@@ -68,8 +144,9 @@ class CivitaiClient:
params['favorites'] = 'true'
r = self._get('/models', params=params, token=token)
if r.status_code != 200:
log.error(f'CivitAI search: code={r.status_code} reason={getattr(r, "reason", "")}')
return CivitSearchResponse()
message = response_message(r)
log.error(f'CivitAI search: code={r.status_code} message="{message}"')
return CivitSearchResponse(error=message)
data = r.json()
if 'items' not in data:
# single model by numeric query — wrap in search response
@@ -82,7 +159,7 @@ class CivitaiClient:
response = CivitSearchResponse.parse_obj(data)
except Exception as e:
log.error(f'CivitAI search parse error: {e}')
return CivitSearchResponse()
return CivitSearchResponse(error='search response could not be parsed')
# /models rejects server-side level filtering and its nsfw boolean leaks
# Mature+ content, so filter on each model's aggregate nsfwLevel here:
# nsfw on keeps every level, nsfw off/unset keeps SFW (None + Soft).
@@ -94,7 +171,7 @@ class CivitaiClient:
def get_model(self, model_id: int, *, token: str | None = None) -> CivitModel | None:
r = self._get(f'/models/{model_id}', token=token)
if r.status_code != 200:
log.error(f'CivitAI get model: id={model_id} code={r.status_code}')
log.error(f'CivitAI get model: id={model_id} code={r.status_code} message="{response_message(r)}"')
return None
try:
return CivitModel.parse_obj(r.json())
@@ -105,7 +182,7 @@ class CivitaiClient:
def get_version(self, version_id: int, *, token: str | None = None) -> CivitVersion | None:
r = self._get(f'/model-versions/{version_id}', token=token)
if r.status_code != 200:
log.error(f'CivitAI get version: id={version_id} code={r.status_code}')
log.error(f'CivitAI get version: id={version_id} code={r.status_code} message="{response_message(r)}"')
return None
try:
return CivitVersion.parse_obj(r.json())
@@ -116,6 +193,8 @@ class CivitaiClient:
def get_version_by_hash(self, hash_str: str, *, token: str | None = None) -> CivitVersion | None:
r = self._get(f'/model-versions/by-hash/{hash_str}', token=token)
if r.status_code != 200:
if r.status_code != 404:
log.error(f'CivitAI get version by hash: hash={hash_str} code={r.status_code} message="{response_message(r)}"')
return None
try:
return CivitVersion.parse_obj(r.json())
@@ -123,6 +202,70 @@ class CivitaiClient:
log.error(f'CivitAI get version by hash parse error: hash={hash_str} {e}')
return None
def get_version_mini(self, version_id: int, *, token: str | None = None) -> CivitVersionMini | None:
r = self._get(f'/model-versions/mini/{version_id}', token=token)
if r.status_code != 200:
log.error(f'CivitAI get version mini: id={version_id} code={r.status_code} message="{response_message(r)}"')
return None
try:
return CivitVersionMini.parse_obj(r.json())
except Exception as e:
log.error(f'CivitAI get version mini parse error: id={version_id} {e}')
return None
def get_version_ids_by_hash(self, hashes: list[str], *, token: str | None = None) -> tuple[list[dict], dict[str, int]]:
"""{modelVersionId, modelId, hash} rows for SHA256 hashes, plus the status code for each hash whose request failed."""
rows, failed = [], {}
for i in range(0, len(hashes), BY_HASH_IDS_LIMIT):
chunk = hashes[i:i + BY_HASH_IDS_LIMIT]
r = self.send('POST', '/model-versions/by-hash/ids', body=chunk, token=token)
if r.status_code != 200:
log.error(f'CivitAI version ids by hash: count={len(chunk)} code={r.status_code} message="{response_message(r)}"')
failed.update(dict.fromkeys(chunk, r.status_code))
continue
try:
rows.extend(r.json())
except Exception as e:
log.error(f'CivitAI version ids by hash parse error: count={len(chunk)} {e}')
failed.update(dict.fromkeys(chunk, 500))
return rows, failed
def get_versions_by_hash(self, hashes: list[str], *, token: str | None = None) -> tuple[list[CivitVersion], dict[str, int]]:
"""Full versions for SHA256 hashes, plus the status code for each hash whose request failed."""
versions, failed = [], {}
for i in range(0, len(hashes), BY_HASH_LIMIT):
chunk = hashes[i:i + BY_HASH_LIMIT]
r = self.send('POST', '/model-versions/by-hash', body=chunk, token=token)
if r.status_code != 200:
log.error(f'CivitAI versions by hash: count={len(chunk)} code={r.status_code} message="{response_message(r)}"')
failed.update(dict.fromkeys(chunk, r.status_code))
continue
try:
versions.extend([CivitVersion.parse_obj(v) for v in r.json()])
except Exception as e:
log.error(f'CivitAI versions by hash parse error: count={len(chunk)} {e}')
failed.update(dict.fromkeys(chunk, 500))
return versions, failed
def get_models_raw(self, model_ids: list[int], *, token: str | None = None) -> tuple[dict[int, dict], dict[int, int]]:
"""Unparsed /models items keyed by id, plus the status code for each id whose request failed."""
models, failed = {}, {}
for i in range(0, len(model_ids), MODEL_IDS_LIMIT):
chunk = model_ids[i:i + MODEL_IDS_LIMIT]
params = {'ids': ','.join(str(m) for m in chunk), 'limit': MODEL_IDS_LIMIT, 'nsfw': 'true'} # ids query drops NSFW models unless nsfw=true
r = self.send('GET', '/models', params=params, token=token)
if r.status_code != 200:
log.error(f'CivitAI models by id: count={len(chunk)} code={r.status_code} message="{response_message(r)}"')
failed.update(dict.fromkeys(chunk, r.status_code))
continue
try:
for item in r.json().get('items', []):
models[item['id']] = item
except Exception as e:
log.error(f'CivitAI models by id parse error: count={len(chunk)} {e}')
failed.update(dict.fromkeys(chunk, 500))
return models, failed
def get_images(self, *, model_version_id: int | None = None, limit: int | None = None, token: str | None = None) -> list[CivitImage]:
params: dict = {}
if model_version_id is not None:
@@ -131,6 +274,7 @@ class CivitaiClient:
params['limit'] = limit
r = self._get('/images', params=params, token=token)
if r.status_code != 200:
log.error(f'CivitAI get images: code={r.status_code} message="{response_message(r)}"')
return []
data = r.json()
items = data.get('items', [])
@@ -152,6 +296,7 @@ class CivitaiClient:
params['limit'] = limit
r = self._get('/images', params=params, token=token)
if r.status_code != 200:
log.error(f'CivitAI get images: code={r.status_code} message="{response_message(r)}"')
return []
data = r.json()
return data.get('items', [])
@@ -166,6 +311,7 @@ class CivitaiClient:
params['page'] = page
r = self._get('/tags', params=params)
if r.status_code != 200:
log.error(f'CivitAI get tags: code={r.status_code} message="{response_message(r)}"')
return CivitTagResponse()
try:
return CivitTagResponse.parse_obj(r.json())
@@ -183,6 +329,7 @@ class CivitaiClient:
params['page'] = page
r = self._get('/creators', params=params)
if r.status_code != 200:
log.error(f'CivitAI get creators: code={r.status_code} message="{response_message(r)}"')
return CivitCreatorResponse()
try:
return CivitCreatorResponse.parse_obj(r.json())
@@ -193,6 +340,8 @@ class CivitaiClient:
def get_me(self, token: str | None = None) -> CivitUserProfile | None:
r = self._get('/me', token=token)
if r.status_code != 200:
if r.status_code != 401:
log.error(f'CivitAI get me: code={r.status_code} message="{response_message(r)}"')
return None
try:
return CivitUserProfile.parse_obj(r.json())
@@ -211,7 +360,7 @@ class CivitaiClient:
"""Civitai enum lists (ModelType, ModelFileType, BaseModel, ActiveBaseModel, BaseModelType). Public endpoint."""
r = self._get('/enums')
if r.status_code != 200:
log.debug(f'CivitAI enums: code={r.status_code}')
log.debug(f'CivitAI enums: code={r.status_code} message="{response_message(r)}"')
return {}
try:
return r.json()
@@ -255,11 +404,10 @@ class CivitaiClient:
if not isinstance(error, dict):
continue
# Parse ZodError: error.message is a JSON-encoded array of issues
import json as _json
issues = error.get('issues', [])
if not issues:
try:
issues = _json.loads(error.get('message', '[]'))
issues = json.loads(error.get('message', '[]'))
except Exception:
issues = []
for issue in issues:
@@ -287,20 +435,15 @@ class CivitaiClient:
break
except Exception as e:
log.debug(f'CivitAI discover options: key={key} {e}')
# Enrich base-model names with github metadata (group/hidden/ecosystem).
# github also serves as the name-list fallback when both /enums and the
# probe came back empty.
github_entries = fetch_github_base_models()
github_index: dict = {entry['name']: entry for entry in github_entries}
if not result['base_models'] and github_entries:
result['base_models'] = [entry['name'] for entry in github_entries]
# hidden marks names in BaseModel but not in ActiveBaseModel, the retired set; an empty ActiveBaseModel hides nothing
active = set(enums.get('ActiveBaseModel', []) or [])
result['base_models_info'] = [
github_index.get(name, {'name': name, 'type': 'image', 'group': '', 'hidden': False})
{'name': name, 'type': 'image', 'group': '', 'hidden': bool(active) and name not in active}
for name in result['base_models']
]
options_cache = result
options_cache_time = now
log.debug(f'CivitAI options: types={len(result["types"])} sort={len(result["sort"])} period={len(result["period"])} base_models={len(result["base_models"])} (enriched={len(github_index)})')
log.debug(f'CivitAI options: types={len(result["types"])} sort={len(result["sort"])} period={len(result["period"])} base_models={len(result["base_models"])} active={len(active)}')
return result
modules/civitai/download_civitai.py
+142 -94
View File
@@ -28,7 +28,7 @@ class DownloadItem:
token: str | None = None
model_id: int = 0
version_id: int = 0
status: str = "queued" # queued | downloading | verifying | completed | failed | cancelled
status: str = "queued" # queued | downloading | completed | failed | cancelled
progress: float = 0.0
bytes_downloaded: int = 0
bytes_total: int = 0
@@ -167,7 +167,7 @@ class DownloadManager:
# Create temp file name from URL hash
url_hash = hashlib.sha256(item.url.encode('utf-8')).hexdigest()[:8]
temp_file = os.path.join(item.folder, f'{url_hash}.tmp')
final_file = os.path.join(item.folder, item.filename)
final_file = os.path.abspath(os.path.join(item.folder, item.filename))
# Check if already exists
if os.path.isfile(final_file):
@@ -195,13 +195,17 @@ class DownloadManager:
item.status = "downloading"
item.bytes_downloaded = starting_pos
digest = hashlib.sha256()
try:
r = shared.req(item.url, headers=headers if headers else None, stream=True)
if r.status_code not in (200, 206):
from modules.civitai.client_civitai import response_message
reason = response_message(r)
item.status = "failed"
item.error = f'HTTP {r.status_code}'
item.error = f'HTTP {r.status_code}: {reason}' if reason else f'HTTP {r.status_code}'
item.completed_at = datetime.now()
log.error(f'CivitAI download refused: id={item.id} file="{item.filename}" code={r.status_code} message="{reason}"')
return
# A text/* response is an error or login page served with HTTP 200,
@@ -211,7 +215,7 @@ class DownloadManager:
item.status = "failed"
item.error = f'invalid content-type: {content_type}'
item.completed_at = datetime.now()
log.warning(f'CivitAI download invalid content-type: id={item.id} content-type="{content_type}"')
log.warning(f'CivitAI download invalid content-type: id={item.id} file="{item.filename}" content-type="{content_type}"')
return
# A 200 reply to a Range request means the server ignored the range
@@ -222,6 +226,10 @@ class DownloadManager:
starting_pos = 0
item.bytes_downloaded = 0
os.truncate(temp_file, 0)
if starting_pos > 0: # a resumed download's digest must include the partial already on disk
with open(temp_file, 'rb') as partial:
for block in iter(lambda: partial.read(1024 * 1024), b''):
digest.update(block)
total_size = int(r.headers.get('content-length', 0))
item.bytes_total = starting_pos + total_size
@@ -255,6 +263,7 @@ class DownloadManager:
return
f.write(chunk)
digest.update(chunk)
written += len(chunk)
item.bytes_downloaded = written
if item.bytes_total > 0:
@@ -270,7 +279,7 @@ class DownloadManager:
item.status = "failed"
item.error = f'incomplete: expected={expected} got={written}'
item.completed_at = datetime.now()
log.warning(f'CivitAI download incomplete: id={item.id} expected={expected} got={written}')
log.warning(f'CivitAI download incomplete: id={item.id} file="{item.filename}" expected={expected} got={written}')
return
elif written < 1024:
try:
@@ -287,30 +296,22 @@ class DownloadManager:
item.status = "failed"
item.error = str(e)
item.completed_at = datetime.now()
log.error(f'CivitAI download error: id={item.id} {e}')
log.error(f'CivitAI download error: id={item.id} file="{item.filename}" {e}')
return
# Hash verification
if item.expected_hash:
item.status = "verifying"
try:
from modules import hashes
computed = hashes.calculate_sha256(temp_file, quiet=True)
if computed.upper() != item.expected_hash.upper():
discard = getattr(shared.opts, 'civitai_discard_hash_mismatch', True)
if discard:
try:
os.remove(temp_file)
except OSError:
pass
item.status = "failed"
item.error = f'hash mismatch: expected={item.expected_hash[:16]}... got={computed[:16]}...'
item.completed_at = datetime.now()
log.error(f'CivitAI download hash mismatch: id={item.id} expected={item.expected_hash[:16]} got={computed[:16]}')
return
log.warning(f'CivitAI download hash mismatch (kept): id={item.id} expected={item.expected_hash[:16]} got={computed[:16]}')
except Exception as e:
log.warning(f'CivitAI download hash check failed: id={item.id} {e}')
computed = digest.hexdigest()
if item.expected_hash and computed != item.expected_hash.lower():
if getattr(shared.opts, 'civitai_discard_hash_mismatch', True):
try:
os.remove(temp_file)
except OSError:
pass
item.status = "failed"
item.error = f'hash mismatch: expected={item.expected_hash[:16]}... got={computed[:16]}...'
item.completed_at = datetime.now()
log.error(f'CivitAI download hash mismatch: id={item.id} expected={item.expected_hash[:16]} got={computed[:16]}')
return
log.warning(f'CivitAI download hash mismatch (kept): id={item.id} expected={item.expected_hash[:16]} got={computed[:16]}')
# Move temp to final
try:
@@ -319,6 +320,7 @@ class DownloadManager:
item.status = "failed"
item.error = f'rename failed: {e}'
item.completed_at = datetime.now()
log.error(f'CivitAI download rename failed: id={item.id} file="{final_file}" {e}')
return
item.status = "completed"
@@ -326,28 +328,25 @@ class DownloadManager:
item.completed_at = datetime.now()
log.info(f'CivitAI download complete: id={item.id} file="{final_file}" size={item.bytes_downloaded}')
# Write verified hash to cache so check-local finds it immediately
if item.expected_hash:
try:
from modules import hashes
model_type_map = {'Checkpoint': 'checkpoint', 'LORA': 'lora', 'TextualInversion': 'embedding', 'VAE': 'vae'}
prefix = model_type_map.get(item.model_type, item.model_type.lower())
name = os.path.splitext(item.filename)[0]
title = f"{prefix}/{name}"
hashes.cache().add_hash(title, os.path.getmtime(final_file), item.expected_hash.lower())
# the declared hash is cached even on a kept mismatch: it is what CivitAI knows the file by
try:
from modules import hashes
from modules.civitai.filemanage_civitai import loader_kind, hash_cache_title
title = hash_cache_title(loader_kind(final_file), final_file)
if title is not None:
hashes.cache().add_hash(title, os.path.getmtime(final_file), (item.expected_hash or computed).lower())
hashes.save_cache()
except Exception:
pass
except Exception as e:
log.warning(f'CivitAI download hash cache: id={item.id} {e}')
# Download metadata and preview
self._fetch_sidecar(item, final_file)
# Refresh model list and extra-networks cache
try:
from modules.sd_models import list_models
list_models()
except Exception:
pass
from modules.civitai.filemanage_civitai import register_download
register_download(final_file)
except Exception as e:
log.warning(f'CivitAI download register: id={item.id} {e}')
try:
from modules.api.loras import _invalidate_extra_networks
_invalidate_extra_networks()
@@ -372,9 +371,9 @@ class DownloadManager:
if version and version.images:
for img in version.images:
if img.url:
code, _size, _note = download_civit_preview(final_file, img.url, meta=img.meta)
code, _size, note = download_civit_preview(final_file, img.url, meta=img.meta)
if code == 200:
log.info(f'CivitAI preview saved: id={item.id}')
log.info(f'CivitAI preview saved: id={item.id} file="{note}"')
break
if code == 304 and backfill_preview_parameters(final_file, img.url, img.meta):
log.info(f'CivitAI preview backfilled: id={item.id}')
@@ -672,7 +671,7 @@ def backfill_preview_parameters(model_path: str, preview_url: str, meta: dict |
ext = os.path.splitext(preview_url)[1].lower()
base = os.path.splitext(model_path)[0]
if ext in VIDEO_PREVIEW_EXTENSIONS:
if not os.path.exists(base + ext):
if not any(os.path.exists(base + e) for e in VIDEO_PREVIEW_EXTENSIONS):
return False
preview_file = base + '.thumb.jpg'
if not os.path.exists(preview_file):
@@ -694,6 +693,13 @@ def backfill_preview_parameters(model_path: str, preview_url: str, meta: dict |
# ---- Legacy compatibility functions ----
def save_civit_meta(model_path: str, data: dict) -> str:
from modules.json_helpers import writefile
fn = os.path.splitext(model_path)[0] + '.json'
writefile(data, filename=fn, mode='w', silent=True)
return fn
def download_civit_meta(model_path: str, model_id):
fn = os.path.splitext(model_path)[0] + '.json'
url = f'https://civitai.com/api/v1/models/{model_id}'
@@ -701,8 +707,7 @@ def download_civit_meta(model_path: str, model_id):
if r.status_code == 200:
try:
data = r.json()
from modules.json_helpers import writefile
writefile(data, filename=fn, mode='w', silent=True)
save_civit_meta(model_path, data)
log.info(f'CivitAI download: id={model_id} url={url} file="{fn}"')
return r.status_code, len(data), ''
except Exception as e:
@@ -713,65 +718,108 @@ def download_civit_meta(model_path: str, model_id):
return r.status_code, '', ''
VIDEO_CONTENT_TYPES = {'video/mp4': '.mp4', 'video/webm': '.webm'}
def transcoded_video_url(preview_url: str) -> str | None:
"""The CDN's H.264 copy of a CivitAI video, or None for any other URL."""
if '/original=true/' not in preview_url:
return None
return preview_url.replace('/original=true/', '/transcode=true,width=450/', 1)
def download_civit_preview(model_path: str, preview_url: str, meta: dict | None = None):
"""Save the preview behind preview_url beside model_path; on 200 the note is the file the UI shows."""
if model_path is None:
return 500, '', ''
ext = os.path.splitext(preview_url)[1]
preview_file = os.path.splitext(model_path)[0] + ext
is_video = preview_file.lower().endswith(VIDEO_PREVIEW_EXTENSIONS)
is_json = preview_file.lower().endswith('.json')
if is_json:
ext = os.path.splitext(preview_url)[1].lower()
base = os.path.splitext(model_path)[0]
if ext == '.json':
log.warning(f'CivitAI download: url="{preview_url}" skip json')
return 500, '', 'expected preview image got json'
if os.path.exists(preview_file):
return 304, '', 'already exists'
r = shared.req(preview_url, stream=True)
total_size = int(r.headers.get('content-length', 0))
block_size = 16384
written = 0
is_video = ext in VIDEO_PREVIEW_EXTENSIONS
if is_video:
if any(os.path.exists(base + e) for e in VIDEO_PREVIEW_EXTENSIONS):
return 304, '', 'already exists'
candidates = [url for url in (transcoded_video_url(preview_url), preview_url) if url] # the original may be AV1, which OpenCV builds without dav1d cannot decode
else:
if os.path.exists(base + ext):
return 304, '', 'already exists'
candidates = [preview_url]
jobid = shared.state.begin('Download CivitAI')
try:
with open(preview_file, 'wb') as f:
for data in r.iter_content(block_size):
written += len(data)
f.write(data)
if written < 1024:
os.remove(preview_file)
return 400, '', 'removed invalid download'
if is_video:
from modules.civitai.video_helper import save_video_frame
save_video_frame(preview_file)
if meta:
thumb_file = os.path.splitext(preview_file)[0] + '.thumb.jpg'
if os.path.exists(thumb_file):
try:
parameters = civitai_meta_to_parameters(meta)
if parameters and embed_preview_parameters(thumb_file, parameters):
log.debug(f'CivitAI preview embed: file="{thumb_file}"')
except Exception as e:
log.debug(f'CivitAI preview embed skipped: file="{thumb_file}" {e}')
else:
from PIL import Image
img = Image.open(preview_file)
log.info(f'CivitAI download: url={preview_url} file="{preview_file}" size={total_size} image={img.size}')
img.close()
for url in candidates:
r = shared.req(url, stream=True)
headers = getattr(r, 'headers', None) or {}
if r.status_code != 200:
log.warning(f'CivitAI preview: url="{url}" code={r.status_code}')
continue
if is_video:
content_type = headers.get('content-type', '').split(';')[0].strip().lower()
file_ext = VIDEO_CONTENT_TYPES.get(content_type)
if file_ext is None:
log.warning(f'CivitAI preview: url="{url}" content-type="{content_type}" not a video')
continue
preview_file = base + file_ext # named by content; the URL says .mp4 for webm originals
else:
preview_file = base + ext
total_size = int(headers.get('content-length', 0))
written = 0
with open(preview_file, 'wb') as f:
for data in r.iter_content(16384):
written += len(data)
f.write(data)
if written < 1024:
os.remove(preview_file)
log.warning(f'CivitAI preview: url="{url}" file="{preview_file}" removed invalid download')
continue
if is_video:
from modules.civitai.video_helper import save_video_frame
thumb_file = base + '.thumb.jpg'
if save_video_frame(preview_file) is None or not os.path.exists(thumb_file):
os.remove(preview_file)
continue
log.info(f'CivitAI download: url={url} file="{preview_file}" size={total_size} thumb="{thumb_file}"')
shown = thumb_file
else:
from PIL import Image
img = Image.open(preview_file)
log.info(f'CivitAI download: url={url} file="{preview_file}" size={total_size} image={img.size}')
img.close()
shown = preview_file
if meta:
try:
parameters = civitai_meta_to_parameters(meta)
if parameters and embed_preview_parameters(preview_file, parameters):
log.debug(f'CivitAI preview embed: file="{preview_file}"')
if parameters and embed_preview_parameters(shown, parameters):
log.debug(f'CivitAI preview embed: file="{shown}"')
except Exception as e:
log.debug(f'CivitAI preview embed skipped: file="{preview_file}" {e}')
log.debug(f'CivitAI preview embed skipped: file="{shown}" {e}')
return 200, str(total_size), shown
return 415, '', 'no usable preview'
except Exception as e:
log.error(f'CivitAI download error: url={preview_url} file="{preview_file}" written={written} {e}')
shared.state.end(jobid)
log.error(f'CivitAI preview error: url={preview_url} file="{base}" {e}')
return 500, '', str(e)
shared.state.end(jobid)
return 200, str(total_size), ''
finally:
shared.state.end(jobid)
def declared_sha256(version_id: int, url: str, filename: str, token: str | None = None) -> str:
"""SHA256 CivitAI declares for the version file behind url, matched by download URL, then by file name."""
if not version_id:
return ''
from modules.civitai.client_civitai import client
version = client.get_version(version_id, token=token)
if version is None:
return ''
for matches in (lambda f: f.download_url == url, lambda f: f.name == filename):
for f in version.files:
if f.hashes.sha256 and matches(f):
return f.hashes.sha256.lower()
return ''
def download_civit_model(model_url: str, model_name: str = '', model_path: str = '', model_type: str = '', token: str | None = None,
base_model: str = '', model_id: int = 0, version_id: int = 0):
base_model: str = '', model_id: int = 0, version_id: int = 0, expected_hash: str = ''):
"""Legacy function — delegates to DownloadManager for non-blocking downloads."""
if not model_url:
log.error('Model download: no url provided')
@@ -791,20 +839,20 @@ def download_civit_model(model_url: str, model_name: str = '', model_path: str =
folder = model_path
else:
folder = os.path.join(paths.models_path, model_path)
expected_hash = expected_hash or declared_sha256(version_id, model_url, model_name, token=token)
item = download_manager.enqueue(
url=model_url,
folder=folder,
filename=model_name or "Unknown",
model_type=model_type,
expected_hash=expected_hash,
token=token,
model_id=model_id,
version_id=version_id,
)
# Wait for completion (legacy blocking behavior)
while item.status in ("queued", "downloading", "verifying"):
while item.status in ("queued", "downloading"):
time.sleep(0.5)
if item.status == "completed" and not item.error:
from modules.sd_models import list_models
list_models()
return os.path.join(item.folder, item.filename)
return None
modules/civitai/filemanage_civitai.py
+150 -3
View File
@@ -5,11 +5,16 @@ from modules.logger import log
# Map CivitAI model types to shared.opts directory settings and fallback subfolder
# names. 'Text Encoder' is a file type, not a model type: versions bundle companion
# files, and clients route those by the file's own type.
# names. 'Text Encoder' is the file type of bundled companion files, 'TextEncoder'
# the model type. Types absent here land in Stable-diffusion.
TYPE_MAP = {
'Checkpoint': ('ckpt_dir', 'Stable-diffusion'),
'Text Encoder': ('te_dir', 'Text-encoder'),
'TextEncoder': ('te_dir', 'Text-encoder'),
'UNet': ('unet_dir', 'UNET'),
'CLIP': ('clip_models_path', 'CLIP'),
'CLIPVision': ('clip_models_path', 'CLIP'),
'Detection': ('yolo_dir', 'yolo'),
'TextualInversion': ('embeddings_dir', 'embeddings'),
'Hypernetwork': ('hypernetwork_dir', 'hypernetworks'),
'AestheticGradient': ('ckpt_dir', 'Stable-diffusion'),
@@ -26,6 +31,9 @@ TYPE_MAP = {
'Other': ('ckpt_dir', 'Stable-diffusion'),
}
# unmapped types already logged; one warning each per session
warned_types: set[str] = set()
# CivitAI has no type for a standalone transformer, so DiT finetunes ship as
# 'Checkpoint' like full models. Bases listed here are full checkpoints and stay
# in Stable-diffusion; any other base is transformer-only in practice and routes
@@ -53,6 +61,9 @@ def get_type_folder(model_type: str, base_model: str = '') -> Path:
return Path(paths.models_path) / custom[model_type]
except Exception as e:
log.warning(f'CivitAI type folder override parse error: {e}')
if model_type not in TYPE_MAP and model_type not in warned_types:
warned_types.add(model_type)
log.warning(f'CivitAI type unmapped: type="{model_type}" folder="Stable-diffusion"')
opt_attr, fallback_dir = TYPE_MAP.get(model_type, ('ckpt_dir', 'Stable-diffusion'))
if model_type == 'Checkpoint' and base_model and not is_full_checkpoint_base(base_model):
opt_attr, fallback_dir = 'unet_dir', 'UNET'
@@ -82,6 +93,142 @@ def iter_type_roots() -> set[Path]:
return {r for r in roots if r.is_dir()}
def path_under(filename: str, root: str | None) -> bool:
if not root:
return False
root = os.path.normcase(os.path.abspath(root)).rstrip(os.sep) + os.sep
return os.path.normcase(os.path.abspath(filename)).startswith(root)
def loader_kind(filename: str) -> str | None:
"""Model loader that lists this file, judged by the folder it is in."""
from modules import shared, paths
ckpt_roots = (getattr(shared.opts, 'ckpt_dir', ''), os.path.join(paths.models_path, 'Stable-diffusion'))
if path_under(filename, getattr(shared.opts, 'vae_dir', '')) or path_under(filename, os.path.join(paths.models_path, 'VAE')):
return 'vae'
if filename.endswith('.vae.safetensors') and any(path_under(filename, root) for root in ckpt_roots):
return 'vae'
if path_under(filename, getattr(shared.opts, 'unet_dir', '')):
return 'unet'
if path_under(filename, getattr(shared.cmd_opts, 'lora_dir', '')):
return 'lora'
if any(path_under(filename, root) for root in ckpt_roots):
return 'checkpoint'
return None
def register_download(filename: str):
"""Add a finished download to its loader's list: one file for lora, a folder scan for unet, vae and checkpoint."""
kind = loader_kind(filename)
if kind == 'lora':
from modules.lora.lora_load import add_network
add_network(filename)
elif kind == 'unet':
from modules.sd_unet import refresh_unet_list
refresh_unet_list()
elif kind == 'vae':
from modules.sd_vae import refresh_vae_list
refresh_vae_list()
elif kind == 'checkpoint':
from modules.sd_models import list_models
list_models()
def hash_cache_title(kind: str | None, filename: str, name: str | None = None) -> str | None:
"""Hash cache key the loader of kind reads for filename, or None when it keeps none."""
from modules import shared, paths
basename = os.path.basename(filename)
stem = os.path.splitext(basename)[0]
if kind == 'lora': # lora_load registers the basename with dots replaced
return 'lora/' + stem.replace('.', '_')
if kind == 'unet': # sd_unet keeps the extension on anything but safetensors
return f"unet/{name or (stem if '.safetensors' in basename else basename)}"
if kind == 'vae':
return f'vae/{os.path.abspath(filename)}'
if kind == 'checkpoint':
if name is None: # CheckpointInfo matches the folder by string prefix, then drops the extension
relname = filename
ckpt_dir = getattr(shared.opts, 'ckpt_dir', '') or ''
model_path = os.path.abspath(os.path.join(paths.models_path, 'Stable-diffusion'))
if ckpt_dir and relname.startswith(ckpt_dir):
relname = os.path.relpath(filename, ckpt_dir)
elif relname.startswith(model_path):
relname = os.path.relpath(filename, model_path)
name = os.path.splitext(relname)[0]
return f'checkpoint/{name}'
return None
hash_cache_pruned = False
def loader_root(kind: str) -> str | None:
"""Folder the loader of kind lists, or None for a kind no loader owns."""
from modules import shared, paths
if kind == 'vae':
return getattr(shared.opts, 'vae_dir', '') or os.path.join(paths.models_path, 'VAE')
if kind == 'unet':
return getattr(shared.opts, 'unet_dir', '')
if kind == 'lora':
return getattr(shared.cmd_opts, 'lora_dir', '')
if kind == 'checkpoint':
return getattr(shared.opts, 'ckpt_dir', '') or os.path.join(paths.models_path, 'Stable-diffusion')
return None
def loader_registry(kind: str) -> dict[str, str] | None:
"""Name to path map of the loader that reads the kind's hash cache keys, or None for a kind no loader owns."""
if kind == 'vae':
from modules.sd_vae import vae_dict
return vae_dict
if kind == 'unet':
from modules.sd_unet import unet_dict
return unet_dict
if kind == 'lora':
from modules.lora.lora_load import available_networks
return {name: entry.filename for name, entry in available_networks.items()}
if kind == 'checkpoint':
from modules.sd_checkpoint import checkpoints_list
return {entry.name: entry.filename for entry in checkpoints_list.values()}
return None
def hash_cache_path(title: str) -> str | None:
"""Path the loader registry holds for a hash cache key, or None when no loaded registry names it."""
kind, _, name = title.partition('/')
if kind == 'vae' and os.path.isabs(name):
return name
registry = loader_registry(kind)
return registry.get(name) if registry else None
def hash_cache_stale(title: str) -> bool:
"""True when the key's loader folder is reachable but the registry no longer lists the file, or lists a path that is gone."""
kind, _, name = title.partition('/')
if kind == 'vae' and os.path.isabs(name):
return os.path.isdir(os.path.dirname(name)) and not os.path.exists(name)
root = loader_root(kind)
if not root or not os.path.isdir(root):
return False
path = (loader_registry(kind) or {}).get(name)
return path is None or not os.path.exists(path)
def prune_hash_cache():
"""Drop hash cache entries for files that are gone, once per process."""
global hash_cache_pruned # pylint: disable=global-statement
if hash_cache_pruned:
return
hash_cache_pruned = True
from modules import hashes
gone = [title for title in list(hashes.cache()) if hash_cache_stale(title)]
for title in gone:
hashes.cache().pop(title, None)
if gone:
hashes.save_cache()
log.info(f'CivitAI hash cache: pruned={len(gone)} entries without files')
def resolve_save_path(model_type: str, model_name: str = "", base_model: str = "",
nsfw: bool = False, creator: str = "", model_id: int = 0,
version_id: int = 0, version_name: str = "") -> Path:
@@ -107,7 +254,7 @@ def resolve_save_path(model_type: str, model_name: str = "", base_model: str = "
for key, value in replacements.items():
subfolder = subfolder.replace(key, value)
# Clean up empty path segments
subfolder = re.sub(r'[/\\]+', os.sep, subfolder)
subfolder = re.sub(r'[/\\]+', lambda _match: os.sep, subfolder) # callable repl: a string repl reads the Windows backslash as an escape
subfolder = subfolder.strip(os.sep)
return base_folder / subfolder
modules/civitai/metadata_civitai.py
+245 -155
View File
@@ -1,8 +1,19 @@
import os
import re
import time
from modules.shared import log, opts, readfile, max_workers
import threading
import concurrent.futures
from modules.shared import log, opts, max_workers, state, cmd_opts
from modules.civitai.client_civitai import client
from modules.civitai.filemanage_civitai import hash_cache_title
GIB = 1024 ** 3
sweep_lock = threading.Lock()
class SweepBusy(Exception):
"""A metadata sweep was started while another one was running."""
class CivitModel:
@@ -24,6 +35,98 @@ class CivitModel:
self.status = 'Not found'
PAGE_KINDS = {'model': 'checkpoint', 'lora': 'lora', 'unet/dit': 'unet', 'vae': 'vae'}
def cache_title(page: str, item: dict) -> str | None:
"""Hash cache key read by the page's own loader, or None when it keeps no cache entry."""
kind = PAGE_KINDS.get(page)
return hash_cache_title(kind, item.get('filename') or '', name=item.get('name') if kind in ('checkpoint', 'unet') else None)
def resolve_sha256(entries: list[tuple[str, dict]], size_limit: int | None = None) -> tuple[dict[str, str], dict[str, str]]:
"""File SHA256 per filename plus a note per unresolved file, from the default hash store or by hashing files below size_limit.
The hashes-addnet store and sshs_model_hash are kohya tensor hashes, which CivitAI never matches.
"""
from modules import hashes
resolved, notes, todo, seen = {}, {}, [], set()
for page, item in entries:
fn = item.get('filename') or ''
if fn in seen or not os.path.isfile(fn):
continue
seen.add(fn)
title = cache_title(page, item)
sha = hashes.sha256_from_cache(fn, title) if title else None
if sha:
resolved[fn] = sha.lower()
elif size_limit is not None and os.path.getsize(fn) >= size_limit:
notes[fn] = f'not hashed: {size_limit // GIB} GiB or larger'
else:
todo.append((fn, title))
if len(todo) == 0:
return resolved, notes
if cmd_opts.no_hashing:
log.warning(f'CivitAI metadata: unhashed={len(todo)} hashing disabled')
notes.update({fn: 'not hashed: hashing disabled' for fn, _title in todo})
return resolved, notes
def hash_file(fn: str) -> str | None:
return None if state.interrupted else hashes.calculate_sha256(fn, quiet=True)
log.info(f'CivitAI metadata: hashing files={len(todo)} size={sum(os.path.getsize(fn) for fn, _title in todo) / GIB:.1f}GB')
jobid = state.begin('CivitAI hash')
state.job_count = len(todo)
cached = 0
try:
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = {executor.submit(hash_file, fn): (fn, title) for fn, title in todo}
for future in concurrent.futures.as_completed(futures):
fn, title = futures[future]
state.job_no += 1
state.textinfo = os.path.basename(fn)
try:
sha = future.result()
except Exception as e:
log.error(f'CivitAI metadata hash: file="{fn}" {e}')
notes[fn] = 'not hashed: unreadable'
continue
if sha is None:
notes[fn] = 'not hashed: interrupted'
continue
resolved[fn] = sha.lower()
if title is not None:
hashes.cache().add_hash(title, os.path.getmtime(fn), resolved[fn])
cached += 1
finally:
if cached > 0:
hashes.save_cache()
state.end(jobid)
return resolved, notes
def apply_update_status(model: CivitModel, sha: str, versions: list[dict], local_hashes: set[str]):
if len(versions) == 0:
return
model.latest = versions[0].get('name', '')
latest_hashes = {str((f.get('hashes') or {}).get('SHA256', '')).upper() for f in versions[0].get('files', [])} - {''}
model.latest_hashes = sorted(latest_hashes)
for ver in versions:
for f in ver.get('files', []):
if str((f.get('hashes') or {}).get('SHA256', '')).upper() != sha.upper():
continue
model.vername = ver.get('name', '')
model.url = f.get('downloadUrl', None)
model.latest_name = f.get('name', '')
if model.vername == model.latest:
model.status = 'Latest version'
elif len(local_hashes & latest_hashes) > 0:
model.status = 'Update downloaded'
else:
model.status = 'Update available'
return
def civit_update_metadata(raw: bool = False):
def create_update_metadata_table(rows: list[CivitModel]):
html = """
@@ -54,132 +157,81 @@ def civit_update_metadata(raw: bool = False):
log.error(f'Model list: row={row} {e}')
return html.format(tbody=tbody)
log.debug('CivitAI update metadata: models')
from modules import ui_extra_networks
from modules.civitai.download_civitai import download_civit_meta
pages = ui_extra_networks.get_pages('Model')
if len(pages) == 0:
return 'CivitAI update metadata: no models found'
page: ui_extra_networks.ExtraNetworksPage = pages[0]
results = []
all_hashes = [(item.get('hash', None) or 'XXXXXXXX').upper()[:8] for item in page.list_items()]
for item in page.list_items():
model = CivitModel(name=item['name'], fn=item['filename'], sha=item.get('hash', None), meta=item.get('metadata', {}))
if model.sha is None or len(model.sha) == 0:
log.debug(f'CivitAI skip search: name="{model.name}" hash=None')
else:
version = client.get_version_by_hash(model.sha)
if version is not None:
model.id = version.model_id
download_civit_meta(model.fn, model.id)
fn = os.path.splitext(item['filename'])[0] + '.json'
model.meta = readfile(fn, silent=True, as_type="dict")
model.name = model.meta.get('name', model.name)
model.versions = len(model.meta.get('modelVersions', []))
time.sleep(0.25) # rate limiting
versions = model.meta.get('modelVersions', [])
if len(versions) > 0:
model.latest = versions[0].get('name', '')
model.latest_hashes.clear()
for v in versions[0].get('files', []):
for h in v.get('hashes', {}).values():
model.latest_hashes.append(h[:8].upper())
for ver in versions:
for f in ver.get('files', []):
for h in f.get('hashes', {}).values():
if h[:8].upper() == model.sha[:8].upper():
model.vername = ver.get('name', '')
model.url = f.get('downloadUrl', None)
model.latest_name = f.get('name', '')
if model.vername == model.latest:
model.status = 'Latest version'
elif any(map(lambda v: v in model.latest_hashes, all_hashes)): # pylint: disable=cell-var-from-loop
model.status = 'Update downloaded'
else:
model.status = 'Update available'
break
results.append(model)
yield results if raw else create_update_metadata_table(results)
yield results if raw else create_update_metadata_table(results)
def atomic_civit_search_metadata(item, results):
from modules.civitai.download_civitai import download_civit_preview, download_civit_meta, backfill_preview_parameters, preview_has_parameters, resolve_preview_file
if item is None:
if not sweep_lock.acquire(blocking=False): # pylint: disable=consider-using-with
log.warning('CivitAI update metadata: another metadata sweep is running')
if raw:
raise SweepBusy('CivitAI metadata sweep already running')
yield 'CivitAI update metadata: another metadata sweep is running'
return
try:
meta = os.path.splitext(item['filename'])[0] + '.json'
except Exception:
return
has_meta = os.path.isfile(meta) and os.stat(meta).st_size > 0
needs_backfill = False
if has_meta and 'missing.png' not in item.get('preview', ''):
actual_preview = resolve_preview_file(item)
if actual_preview and not preview_has_parameters(actual_preview):
needs_backfill = True
if ('missing.png' in item['preview'] or not has_meta or needs_backfill) and os.path.isfile(item['filename']):
sha = item.get('hash', None)
found = False
result = {
'id': '',
'name': item['name'],
'type': '',
'hash': '',
'code': '',
'size': '',
'note': '',
}
if sha is not None and len(sha) > 0:
version = client.get_version_by_hash(sha)
result['hash'] = sha
if version is not None:
result['code'] = 200
result['code'], result['size'], result['note'] = download_civit_meta(item['filename'], version.model_id)
result['id'] = version.model_id
result['type'] = 'metadata'
# Create a new dict for each append to avoid mutation bugs
results.append(dict(result))
for img in version.images:
if img.url:
code, size, note = download_civit_preview(item['filename'], img.url, meta=img.meta)
if code == 200:
results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'})
found = True
break
if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta):
results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'})
found = True
break
else:
result['code'] = 404
time.sleep(0.25) # rate limiting
if not found and os.stat(item['filename']).st_size < (1024 * 1024 * 1024):
from modules import hashes
sha = hashes.calculate_sha256(item['filename'], quiet=True)[:10]
version = client.get_version_by_hash(sha)
result['hash'] = sha
if version is not None:
result['code'] = 200
result['code'], result['size'], result['note'] = download_civit_meta(item['filename'], version.model_id)
result['id'] = version.model_id
result['type'] = 'metadata'
results.append(dict(result))
for img in version.images:
if img.url:
code, size, note = download_civit_preview(item['filename'], img.url, meta=img.meta)
if code == 200:
results.append({**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'})
found = True
break
if code == 304 and backfill_preview_parameters(item['filename'], img.url, img.meta):
results.append({**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'})
found = True
break
else:
result['code'] = 404
time.sleep(0.25) # rate limiting
if not found:
results.append(dict(result))
log.debug('CivitAI update metadata: models')
from modules import ui_extra_networks
from modules.civitai.download_civitai import save_civit_meta
pages = ui_extra_networks.get_pages('Model')
if len(pages) == 0:
yield [] if raw else 'CivitAI update metadata: no models found'
return
items = [item for item in pages[0].list_items() if item is not None]
shas, notes = resolve_sha256([('model', item) for item in items])
rows, failed = client.get_version_ids_by_hash(sorted(set(shas.values())))
by_hash = {}
for row in rows: # one hash can match several versions; GET /by-hash/{hash} returns the first
by_hash.setdefault(str(row.get('hash', '')).lower(), row)
metas, failed_models = client.get_models_raw(sorted({row['modelId'] for row in by_hash.values() if row.get('modelId')}))
local_hashes = {sha.upper() for sha in shas.values()}
results = []
for item in items:
fn = item['filename']
sha = shas.get(fn)
model = CivitModel(name=item['name'], fn=fn, sha=sha[:10] if sha else item.get('hash', None))
if sha is None:
model.status = 'Not hashed' if fn in notes else 'Not found'
elif sha in failed:
model.status = 'Lookup failed'
elif sha in by_hash:
model.id = by_hash[sha]['modelId']
meta = metas.get(model.id)
if meta is None:
model.status = 'Lookup failed' if model.id in failed_models else 'Not found'
else:
meta_fn = save_civit_meta(fn, meta)
log.info(f'CivitAI download: id={model.id} file="{meta_fn}"')
model.meta = meta
model.name = meta.get('name', model.name)
model.versions = len(meta.get('modelVersions', []))
apply_update_status(model, sha, meta.get('modelVersions', []), local_hashes)
results.append(model)
yield results if raw else create_update_metadata_table(results)
yield results if raw else create_update_metadata_table(results)
finally:
sweep_lock.release()
def needs_metadata(item: dict) -> bool:
"""True when the item's file exists and it lacks a sidecar, a preview, or preview parameters."""
from modules.civitai.download_civitai import preview_has_parameters, resolve_preview_file
filename = item.get('filename') or ''
if not os.path.isfile(filename):
return False
meta = os.path.splitext(filename)[0] + '.json'
if 'missing.png' in (item.get('preview') or '') or not (os.path.isfile(meta) and os.stat(meta).st_size > 0):
return True
actual_preview = resolve_preview_file(item)
return bool(actual_preview) and not preview_has_parameters(actual_preview)
def download_previews(fn: str, version, result: dict) -> list[dict]:
"""Preview row for the first version image that downloads or gains embedded parameters."""
from modules.civitai.download_civitai import download_civit_preview, backfill_preview_parameters
for img in version.images:
if not img.url:
continue
code, size, note = download_civit_preview(fn, img.url, meta=img.meta)
if code == 200:
return [{**result, 'code': code, 'size': size, 'note': note, 'type': 'preview'}]
if code == 304 and backfill_preview_parameters(fn, img.url, img.meta):
return [{**result, 'code': 200, 'size': '', 'note': 'metadata embedded', 'type': 'preview'}]
return []
def civit_search_metadata(title: str | None = None, raw: bool = False):
@@ -208,36 +260,74 @@ def civit_search_metadata(title: str | None = None, raw: bool = False):
log.error(f'Model list: row={row} {e}')
return html.format(tbody=tbody)
from modules.ui_extra_networks import get_pages
results = []
scanned, skipped = 0, 0
t0 = time.time()
candidates = []
re_skip = [r.strip() for r in opts.extra_networks_scan_skip.split(',') if len(r.strip()) > 0]
for page in get_pages():
if isinstance(title, str):
if page.title.lower() != title.lower():
if not sweep_lock.acquire(blocking=False): # pylint: disable=consider-using-with
log.warning('CivitAI search metadata: another metadata sweep is running')
if raw:
raise SweepBusy('CivitAI metadata sweep already running')
yield 'CivitAI search metadata: another metadata sweep is running'
return
try:
from modules.ui_extra_networks import get_pages
from modules.civitai.download_civitai import save_civit_meta
results = []
scanned, skipped = 0, 0
t0 = time.time()
entries = []
re_skip = [r.strip() for r in opts.extra_networks_scan_skip.split(',') if len(r.strip()) > 0]
for page in get_pages():
if isinstance(title, str) and page.title.lower() != title.lower():
continue
if page.name in ('style', 'wildcards'):
continue
for item in page.list_items():
if item is None:
if page.name in ('style', 'wildcards'):
continue
if any(re.search(re_str, item.get('name', '') + item.get('filename', '')) for re_str in re_skip):
skipped += 1
for item in page.list_items():
if item is None:
continue
if any(re.search(re_str, item.get('name', '') + item.get('filename', '')) for re_str in re_skip):
skipped += 1
continue
scanned += 1
entries.append((page.name, item))
log.debug(f'CivitAI search metadata: type={title if isinstance(title, str) else "all"} workers={max_workers} skip={len(re_skip)} items={len(entries)}')
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
flags = list(executor.map(lambda entry: needs_metadata(entry[1]), entries))
entries = [entry for entry, flag in zip(entries, flags) if flag]
shas, notes = resolve_sha256(entries, size_limit=GIB)
versions, failed = client.get_versions_by_hash(sorted(set(shas.values())))
by_hash = {}
for version in versions: # one hash can match several versions; GET /by-hash/{hash} returns the first
for f in version.files:
if f.hashes.sha256:
by_hash.setdefault(f.hashes.sha256.lower(), version)
matched = {}
for _page, item in entries:
fn = item['filename']
sha = shas.get(fn)
result = {'id': '', 'name': item['name'], 'type': '', 'hash': sha[:10] if sha else '', 'code': '', 'size': '', 'note': notes.get(fn, '')}
if sha is None:
results.append(result)
elif sha in failed:
results.append({**result, 'code': failed[sha], 'note': 'lookup failed'})
elif sha not in by_hash:
results.append({**result, 'code': 404})
else:
matched[fn] = (by_hash[sha], {**result, 'id': by_hash[sha].model_id})
metas, failed_models = client.get_models_raw(sorted({version.model_id for version, _result in matched.values()}))
for fn, (version, result) in matched.items():
meta = metas.get(version.model_id)
if meta is None:
results.append({**result, 'type': 'metadata', 'code': failed_models.get(version.model_id, 404)})
continue
scanned += 1
candidates.append(item)
log.debug(f'CivitAI search metadata: type={title if isinstance(title, str) else "all"} workers={max_workers} skip={len(re_skip)} items={len(candidates)}')
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
future_items = {}
for candidate in candidates:
future_items[executor.submit(atomic_civit_search_metadata, candidate, results)] = candidate
for future in concurrent.futures.as_completed(future_items):
future.result()
yield results if raw else create_search_metadata_table(results)
t1 = time.time()
log.debug(f'CivitAI search metadata: scanned={scanned} skipped={skipped} time={t1 - t0:.2f}')
yield results if raw else create_search_metadata_table(results)
meta_fn = save_civit_meta(fn, meta)
log.info(f'CivitAI download: id={version.model_id} file="{meta_fn}"')
results.append({**result, 'type': 'metadata', 'code': 200, 'size': len(meta)})
yield results if raw else create_search_metadata_table(results)
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
futures = [executor.submit(download_previews, fn, version, result) for fn, (version, result) in matched.items()]
for future in concurrent.futures.as_completed(futures):
results.extend(future.result())
yield results if raw else create_search_metadata_table(results)
t1 = time.time()
log.debug(f'CivitAI search metadata: scanned={scanned} skipped={skipped} pending={len(entries)} hashed={len(shas)} matched={len(matched)} time={t1 - t0:.2f}')
yield results if raw else create_search_metadata_table(results)
finally:
sweep_lock.release()
modules/civitai/models_civitai.py
+43 -9
View File
@@ -23,6 +23,7 @@ class CivitFileHashes(BaseModel):
autov3: str | None = Field(None, alias="AutoV3")
crc32: str | None = Field(None, alias="CRC32")
blake3: str | None = Field(None, alias="BLAKE3")
sha256_12: str | None = Field(None, alias="SHA256_12")
class CivitFileMetadata(BaseModel):
@@ -51,15 +52,14 @@ class CivitFile(BaseModel):
class CivitStats(BaseModel):
# counts CivitAI does not report arrive as null, not zero
class Config:
allow_population_by_field_name = True
download_count: int = Field(0, alias="downloadCount")
favorite_count: int = Field(0, alias="favoriteCount")
thumb_up_count: int = Field(0, alias="thumbsUpCount")
thumb_down_count: int = Field(0, alias="thumbsDownCount")
comment_count: int = Field(0, alias="commentCount")
rating_count: int = Field(0, alias="ratingCount")
rating: float = 0
download_count: int | None = Field(0, alias="downloadCount")
thumb_up_count: int | None = Field(0, alias="thumbsUpCount")
thumb_down_count: int | None = Field(0, alias="thumbsDownCount")
comment_count: int | None = Field(0, alias="commentCount")
tipped_amount_count: int | None = Field(0, alias="tippedAmountCount")
class CivitVersion(BaseModel):
@@ -89,6 +89,36 @@ class CivitVersion(BaseModel):
return "Unknown" if v in (None, "") else v
class CivitVersionMini(BaseModel):
# primary file flattened onto the version plus permission flags; earlyAccessEndsAt and freeTrialLimit exist only during early access
class Config:
allow_population_by_field_name = True
air: str = ""
version_name: str = Field("", alias="versionName")
model_name: str = Field("", alias="modelName")
user_id: int = Field(0, alias="userId")
base_model: str = Field("Unknown", alias="baseModel")
availability: str = "Unknown"
published_at: str | None = Field(None, alias="publishedAt")
size: float = 0
file_type: str = Field("", alias="fileType")
file_name: str = Field("", alias="fileName")
format: str = ""
hashes: CivitFileHashes = Field(default_factory=CivitFileHashes)
download_urls: list[str] = Field(default_factory=list, alias="downloadUrls")
can_generate: bool = Field(False, alias="canGenerate")
is_featured: bool = Field(False, alias="isFeatured")
require_auth: bool = Field(False, alias="requireAuth")
check_permission: bool = Field(False, alias="checkPermission")
additional_resource_charge: bool = Field(False, alias="additionalResourceCharge")
payout_enabled: bool = Field(False, alias="payoutEnabled")
minor: bool = False
sfw_only: bool = Field(False, alias="sfwOnly")
fees: list = Field(default_factory=list)
early_access_ends_at: str | None = Field(None, alias="earlyAccessEndsAt")
free_trial_limit: int | None = Field(None, alias="freeTrialLimit")
class CivitCreator(BaseModel):
class Config:
allow_population_by_field_name = True
@@ -158,6 +188,7 @@ class CivitSearchResponse(BaseModel):
items: list[CivitModel] = Field(default_factory=list)
metadata: CivitSearchMetadata = Field(default_factory=CivitSearchMetadata)
request_url: str | None = Field(None, alias="requestUrl")
error: str | None = None # server or parse failure text; items is empty when set
class CivitTag(BaseModel):
@@ -192,9 +223,12 @@ class CivitCreatorResponse(BaseModel):
class CivitUserProfile(BaseModel):
# tier is omitted for non-members; email, emailVerified and tokenScope are left unmodelled to keep the address out of the API response
class Config:
allow_population_by_field_name = True
id: int = 0
username: str = ""
image: str | None = None
profile_picture: str | None = Field(None, alias="profilePicture")
tier: str | None = None
status: str | None = None
is_member: bool = Field(False, alias="isMember")
subscriptions: list = Field(default_factory=list)
modules/civitai/search_civitai.py
+16 -16
View File
@@ -1,14 +1,11 @@
import re
import time
from html import escape
from installer import log
from modules.civitai.client_civitai import client
from modules.civitai.models_civitai import CivitModel, CivitSearchResponse
# Hardcoded fallback list — used by Gradio UI if discover_options() fails
base_models = ['', 'AuraFlow', 'Chroma', 'CogVideoX', 'Flux.1 S', 'Flux.1 D', 'Flux.1 Krea', 'Flux.1 Kontext', 'Flux.2 D', 'HiDream', 'Hunyuan 1', 'Hunyuan Video', 'Illustrious', 'Kolors', 'LTXV', 'Lumina', 'Mochi', 'NoobAI', 'PixArt a', 'PixArt E', 'Pony', 'Pony V7', 'Qwen', 'SD 1.4', 'SD 1.5', 'SD 1.5 LCM', 'SD 1.5 Hyper', 'SD 2.0', 'SD 2.1', 'SDXL 1.0', 'SDXL Lightning', 'SDXL Hyper', 'Wan Video 1.3B t2v', 'Wan Video 14B t2v', 'Wan Video 14B i2v 480p', 'Wan Video 14B i2v 720p', 'Wan Video 2.2 TI2V-5B', 'Wan Video 2.2 I2V-A14B', 'Wan Video 2.2 T2V-A14B', 'Wan Video 2.5 T2V', 'Wan Video 2.5 I2V', 'ZImageTurbo', 'Other']
def search_civitai(
query: str,
tag: str = '',
@@ -20,10 +17,10 @@ def search_civitai(
base: str = '',
token: str | None = None,
exact: bool = True,
) -> list[CivitModel]:
) -> CivitSearchResponse:
if not query and not tag and not sort:
log.error('CivitAI: no search criteria provided')
return []
return CivitSearchResponse(error='no search criteria provided')
t0 = time.time()
@@ -39,8 +36,8 @@ def search_civitai(
if model:
t1 = time.time()
log.info(f'CivitAI result: id={query} time={t1 - t0:.2f}')
return [model]
return []
return CivitSearchResponse(items=[model])
return CivitSearchResponse(error=f'model {query} not found')
response: CivitSearchResponse = client.search_models(
query=query,
@@ -65,17 +62,20 @@ def search_civitai(
if any(q_lower in name for name in names):
exact_models.append(model)
result = exact_models if exact_models else all_models
response.items = exact_models if exact_models else all_models
t1 = time.time()
log.info(f'CivitAI result: exact={len(exact_models)} total={len(all_models)} time={t1 - t0:.2f}')
return result
return response
def create_model_cards(all_models: list[CivitModel]) -> str:
details = """
<div id="model-details">
</div>
"""
def create_model_cards(response: CivitSearchResponse) -> str:
if response.error:
notice = f'CivitAI: {escape(response.error)}'
elif not response.items:
notice = 'No models found'
else:
notice = ''
details = f'<div id="model-details">{notice}</div>'
cards = """
<div id="model-cards" class="extra-network-cards">
{cards}
@@ -89,7 +89,7 @@ def create_model_cards(all_models: list[CivitModel]) -> str:
</div>
"""
all_cards = ''
for model in all_models:
for model in response.items:
previews = []
for version in model.versions:
for image in version.images:
modules/control/proc/hed.py
+2 -2
View File
@@ -102,8 +102,8 @@ class HEDdetector:
if scribble:
detected_map = nms(detected_map, 127, 3.0)
detected_map = cv2.GaussianBlur(detected_map, (0, 0), 3.0)
detected_map[detected_map > 4] = 255
detected_map[detected_map < 255] = 0
detected_map[detected_map > 4] = 255 # pylint: disable=unsupported-assignment-operation
detected_map[detected_map < 255] = 0 # pylint: disable=unsupported-assignment-operation
if opts.control_move_processor:
self.model.to('cpu')
if output_type == "pil":
modules/control/run.py
+11 -16
View File
@@ -350,9 +350,8 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
styles: list[str] | None = None,
steps: int = 20, sampler_index: int | None = None,
seed: int = -1, subseed: int = -1, subseed_strength: float = 0, seed_resize_from_h: int = -1, seed_resize_from_w: int = -1,
guidance_name: str = 'Default', guidance_scale: float = 6.0, guidance_rescale: float = 0.0, guidance_start: float = 0.0, guidance_stop: float = 1.0,
cfg_scale: float = 6.0, clip_skip: float = 1.0, cfg_image: float = 6.0, cfg_rescale: float = 0.7, cfg_true: float = 0.0, cfg_adaptive: float = 0.5, cfg_end: float = 1.0,
vae_type: str = 'Full', tiling: bool = False, hidiffusion: bool = False,
cfg_name: str = 'Default', cfg_scale: float = 6.0, cfg_image: float = 6.0, cfg_rescale: float = 0.0, cfg_start: float = 0.0, cfg_stop: float = 1.0, cfg_true: float = 0.0, cfg_adaptive: float = 0.5,
clip_skip: float = 1.0, vae_type: str = 'Full', tiling: bool = False, hidiffusion: bool = False,
detailer_enabled: bool = False, detailer_prompt: str = '', detailer_negative: str = '', detailer_steps: int = 10, detailer_strength: float = 0.3, detailer_resolution: int = 1024, detailer_classes: str = '',
hdr_mode: int = 0, hdr_brightness: float = 0, hdr_color: float = 0, hdr_sharpen: float = 0, hdr_clamp: bool = False, hdr_boundary: float = 4.0, hdr_threshold: float = 0.95,
hdr_maximize: bool = False, hdr_max_center: float = 0.6, hdr_max_boundary: float = 1.0, hdr_color_picker: str | None = None, hdr_tint_ratio: float = 0, hdr_apply_hires: bool = True,
@@ -403,7 +402,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
sequential_seed: bool | None = None,
# prompt/attention overrides
prompt_attention: str | None = None, prompt_mean_norm: bool | None = None, diffusers_zeros_prompt_pad: bool | None = None,
te_pooled_embeds: bool | None = None, lora_apply_te: bool | None = None, te_complex_human_instruction: str | None = None, te_use_mask: bool | None = None,
te_pooled_embeds: bool | None = None, te_complex_human_instruction: str | None = None, te_use_mask: bool | None = None,
# generation modifier overrides (hijack)
freeu_enabled: bool | None = None, freeu_b1: float | None = None, freeu_b2: float | None = None, freeu_s1: float | None = None, freeu_s2: float | None = None,
hypertile_unet_enabled: bool | None = None, hypertile_hires_only: bool | None = None, hypertile_unet_tile: int | None = None, hypertile_unet_min_tile: int | None = None,
@@ -445,7 +444,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
if sampler_index is None:
log.warning('Sampler: invalid')
sampler_index = 0
if hr_sampler_index is None:
if hr_sampler_index is None or hr_sampler_index == 'Same as primary':
hr_sampler_index = sampler_index
if isinstance(extra, list):
extra = create_override_settings_dict(extra)
@@ -465,21 +464,17 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
seed_resize_from_w = seed_resize_from_w,
denoising_strength = denoising_strength,
skip_processing = skip_processing,
# modular guidance
guidance_name = guidance_name,
guidance_scale = guidance_scale,
guidance_rescale = guidance_rescale,
guidance_start = guidance_start,
guidance_stop = guidance_stop,
# legacy guidance
# guidance
cfg_name = cfg_name,
cfg_scale = cfg_scale,
cfg_end = cfg_end,
clip_skip = clip_skip,
cfg_image = cfg_image,
cfg_rescale = cfg_rescale,
cfg_start = cfg_start,
cfg_stop = cfg_stop,
cfg_true = cfg_true,
cfg_adaptive = cfg_adaptive,
# advanced
clip_skip = clip_skip,
vae_type = vae_type,
tiling = tiling,
hidiffusion = hidiffusion,
@@ -586,7 +581,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
# prompt/attention overrides
prompt_attention=prompt_attention, prompt_mean_norm=prompt_mean_norm,
diffusers_zeros_prompt_pad=diffusers_zeros_prompt_pad, te_pooled_embeds=te_pooled_embeds,
lora_apply_te=lora_apply_te, te_complex_human_instruction=te_complex_human_instruction, te_use_mask=te_use_mask,
te_complex_human_instruction=te_complex_human_instruction, te_use_mask=te_use_mask,
# generation modifier overrides (hijack)
freeu_enabled=freeu_enabled, freeu_b1=freeu_b1, freeu_b2=freeu_b2, freeu_s1=freeu_s1, freeu_s2=freeu_s2,
hypertile_unet_enabled=hypertile_unet_enabled, hypertile_hires_only=hypertile_hires_only,
@@ -852,7 +847,7 @@ def control_run(state: str = '', # pylint: disable=keyword-arg-before-vararg
debug_log(f'Ready: {image_txt}')
html_txt = f'<p>Ready {image_txt}</p>' if image_txt != '' else ''
if len(info_txt) > 0:
if (info_txt is not None) and (len(info_txt) > 0):
html_txt = html_txt + infotext_to_html(info_txt[0])
result = (output_images, blended_image, html_txt, output_filename)
if is_generator:
modules/detailer/detailer.py
+33 -6
View File
@@ -1,4 +1,5 @@
import re
import time
from copy import copy
import numpy as np
import gradio as gr
@@ -220,6 +221,19 @@ class Detailer():
matched_negative_classes = set()
for i, model_val in enumerate(models):
if shared.state.skipped:
shared.state.skipped = False
continue
if shared.state.interrupted:
break
while shared.state.paused:
log.debug('Detail paused')
if shared.state.interrupted:
break
if shared.state.skipped:
continue
time.sleep(0.1)
if ':' in model_val:
model_name, model_args = model_val.split(':', 1)
else:
@@ -309,10 +323,23 @@ class Detailer():
resolved_prompts = assign_prompts(prompt, items)
resolved_negatives = assign_prompts(negative, items)
for j, item in enumerate(items):
if shared.state.skipped:
shared.state.skipped = False
continue
if shared.state.interrupted:
break
while shared.state.paused:
log.debug('Detail paused')
if shared.state.interrupted:
break
if shared.state.skipped:
continue
time.sleep(0.1)
if item.mask is None:
continue
pc.keep_prompts = True
shared.sd_model.fail_on_switch_error = True
pc.keep_prompts = True
pc.prompt = resolved_prompts[j]
pc.negative_prompt = resolved_negatives[j]
pc.prompts = [pc.prompt]
@@ -321,7 +348,7 @@ class Detailer():
pc.disable_extra_networks = True # disable processing_diffusers from handling network activation since its handled here
network_same = len(p.network_data.values()) == len(pc.network_data.values()) and all(x == y for x, y in zip(p.network_data.values(), pc.network_data.values()))
if not network_same:
extra_networks.activate_filtered(pc, pc.network_data)
extra_networks.activate(pc, pc.network_data)
log.debug(f'Detail: model="{i+1}:{name}" item={j+1}/{len(items)} box={item.box} label="{item.label}" score={item.score:.2f} seg={detailer_opt(p, "detailer_segmentation")} network={network_same} prompt="{pc.prompt}"')
pc.init_images = [image]
pc.image_mask = [item.mask]
@@ -510,12 +537,12 @@ class Detailer():
renoise_end = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, label='Renoise end', value=shared.opts.detailer_sigma_adjust_max, elem_id=f"{tab}_detailer_renoise_end")
sampler_block = None
if tab == 'extras': # fold the standalone sampler settings into the detailer accordion; values applied per-job in make_processing, never global opts
from modules import sd_samplers
sd_samplers.set_samplers()
sampler_choices = [s.name for s in sd_samplers.visible_samplers() if s.name != 'Same as primary']
from modules import ui_sections
sampler_choices, default_value, filtered = ui_sections.sampler_choices()
with gr.Accordion('Sampler', open=False, elem_id=f"{tab}_detailer_sampler_accordion", elem_classes=["small-accordion"]):
with gr.Row():
d_sampler = gr.Dropdown(label='Sampling method', choices=sampler_choices, value='Default', elem_id=f"{tab}_detailer_sampler")
ui_sections.create_filter_indicator(tab, 'Sampler', filtered)
d_sampler = gr.Dropdown(label='Sampling method', choices=sampler_choices, value=default_value, type='value', elem_id=f"{tab}_detailer_sampler")
d_prediction = gr.Dropdown(label='Prediction method', choices=['default', 'epsilon', 'sample', 'v_prediction', 'flow_prediction'], value='default', elem_id=f"{tab}_detailer_prediction")
with gr.Row():
d_shift = gr.Slider(label='Flow shift', minimum=0, maximum=10, step=0.1, value=shared.opts.schedulers_shift, elem_id=f"{tab}_detailer_shift")
modules/devices.py
+38 -43
View File
@@ -493,65 +493,60 @@ def override_ipex_math():
log.warning(f'Torch ipex: {e}')
def report_attention():
from importlib.metadata import version
try:
flash = version('flash-attn')
except Exception:
flash = False
try:
sage = version('sageattention')
except Exception:
sage = False
try:
xformers = version('xformers')
except Exception:
xformers = False
try:
kernels = version('kernels')
except Exception:
kernels = False
from diffusers.models import attention_dispatch as a
try:
import sdnq
sdnq_ver = sdnq.__version__
except Exception:
sdnq_ver = False
# log.debug(f'Attention available: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN} kernels={a.is_kernels_available()} sdnq=True') # pylint: disable=protected-access
log.debug(f'Attention available: sdnq={sdnq_ver} flash={flash} sage={sage} flex={a._CAN_USE_FLEX_ATTN} xformers={xformers} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} kernels={kernels}') # pylint: disable=protected-access
def set_sdpa_params():
try:
global sdpa_original # pylint: disable=global-statement
report = sdpa_original is None
try:
global sdpa_original # pylint: disable=global-statement
if sdpa_original is not None:
torch.nn.functional.scaled_dot_product_attention = sdpa_original
else:
sdpa_original = torch.nn.functional.scaled_dot_product_attention
except Exception as err:
log.warning(f'Torch attention: type="sdpa" {err}')
log.warning(f'Attention: type="SDPA" {err}')
try:
options = {}
torch.backends.cuda.enable_flash_sdp('Flash' in opts.sdp_options or 'Flash attention' in opts.sdp_options)
torch.backends.cuda.enable_mem_efficient_sdp('Memory' in opts.sdp_options or 'Memory attention' in opts.sdp_options)
torch.backends.cuda.enable_math_sdp('Math' in opts.sdp_options or 'Math attention' in opts.sdp_options)
if hasattr(torch.backends.cuda, "allow_fp16_bf16_reduction_math_sdp"): # only valid for torch >= 2.5
options['math'] = 'fp16/bf16'
torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)
torch_info.set(attention="sdpa")
log.debug(f'Torch attention: type="sdpa" kernels={opts.sdp_options} overrides={opts.sdp_overrides}')
log.debug(f'Attention: type="SDPA" kernels={opts.sdp_options} options={options}')
except Exception as err:
log.warning(f'Torch attention: type="sdpa" {err}')
# Stack hijcaks in reverse order. This gives priority to the last added hijack.
# If the last hijack is not compatible, it will use the one before it and so on.
if 'Dynamic attention' in opts.sdp_overrides:
global sdpa_pre_dyanmic_atten # pylint: disable=global-statement
sdpa_pre_dyanmic_atten = attention.set_dynamic_attention()
if 'Flex attention' in opts.sdp_overrides:
attention.set_flex_attention()
if 'Triton Flash attention' in opts.sdp_overrides:
attention.set_triton_flash_attention(backend)
if 'Flash attention' in opts.sdp_overrides:
attention.set_ck_flash_attention(backend, device)
if 'Sage attention' in opts.sdp_overrides:
attention.set_sage_attention(backend, device)
if 'SDNQ attention' in opts.sdp_overrides:
attention.set_sdnq_attention()
from importlib.metadata import version
try:
flash = version('flash-attn')
except Exception:
flash = False
try:
sage = version('sageattention')
except Exception:
sage = False
if flash or sage:
log.debug(f'Torch attention installed: flashattn={flash} sageattention={sage}')
from diffusers.models import attention_dispatch as a
log.debug(f'Torch attention available: flash={a._CAN_USE_FLASH_ATTN} flash3={a._CAN_USE_FLASH_ATTN_3} sage={a._CAN_USE_SAGE_ATTN} flex={a._CAN_USE_FLEX_ATTN} npu={a._CAN_USE_NPU_ATTN} xla={a._CAN_USE_XLA_ATTN} xformers={a._CAN_USE_XFORMERS_ATTN} kernels={a.is_kernels_available()} sdnq=True') # pylint: disable=protected-access
log.warning(f'Attention: type="SDPA" {err}')
attention.install_router([opts.cross_attention_optimization], attention.Platform(backend=backend, device=device), sdpa_original)
if report:
report_attention()
except Exception as e:
log.warning(f'Torch SDPA: {e}')
modules/extensions.py
+12 -23
View File
@@ -73,22 +73,11 @@ def temp_disable_extensions():
'sd-extension-nudenet',
'sd-extension-promptgen',
]
disable_themes = [
'sd-webui-lobe-theme',
'cozy-nest',
'sdnext-modernui',
]
disabled = []
if shared.cmd_opts.theme is not None:
theme_name = shared.cmd_opts.theme
else:
theme_name = f'{shared.opts.theme_type.lower()}/{shared.opts.gradio_theme}'
if theme_name == 'lobe':
disable_themes.remove('sd-webui-lobe-theme')
elif theme_name == 'cozy-nest' or theme_name == 'cozy':
disable_themes.remove('cozy-nest')
elif '/' not in theme_name: # set default themes per type
if theme_name == 'standard' or theme_name == 'default':
theme_type = shared.cmd_opts.theme if shared.cmd_opts.theme is not None else shared.opts.theme_type
theme_name = f'{theme_type.lower()}/{shared.opts.gradio_theme}'
if '/' not in theme_name: # set default themes per type
if theme_name == 'standard':
theme_name = 'standard/black-teal'
if theme_name == 'modern':
theme_name = 'modern/Default'
@@ -97,24 +86,22 @@ def temp_disable_extensions():
if theme_name == 'huggingface':
theme_name = 'huggingface/blaaa'
if theme_name.lower().startswith('standard') or theme_name.lower().startswith('default'):
if theme_name.lower().startswith('standard'):
shared.opts.data['theme_type'] = 'Standard'
shared.opts.data['gradio_theme'] = theme_name[9:]
disabled.append('sdnext-modernui')
elif theme_name.lower().startswith('modern'):
shared.opts.data['theme_type'] = 'Modern'
shared.opts.data['gradio_theme'] = theme_name[7:]
disable_themes.remove('sdnext-modernui')
elif theme_name.lower().startswith('huggingface') or theme_name.lower().startswith('gradio') or theme_name.lower().startswith('none'):
shared.opts.data['theme_type'] = 'None'
shared.opts.data['gradio_theme'] = theme_name
disabled.append('sdnext-modernui')
else:
log.error(f'UI theme invalid: theme="{theme_name}" available={["standard/*", "modern/*", "none/*"]} fallback="standard/black-teal"')
shared.opts.data['theme_type'] = 'Standard'
shared.opts.data['gradio_theme'] = 'black-teal'
shared.opts.data['theme_type'] = 'Modern'
shared.opts.data['gradio_theme'] = 'Default'
for ext in disable_themes:
if ext.lower() not in shared.opts.disabled_extensions:
disabled.append(ext)
if shared.cmd_opts.safe:
for ext in disable_safe:
if ext.lower() not in shared.opts.disabled_extensions:
@@ -271,4 +258,6 @@ def list_extensions():
enabled = dirname.lower() not in disabled_extensions
extension = Extension(name=dirname, path=path, enabled=enabled, is_builtin=is_builtin)
extensions.append(extension)
log.debug(f'Extensions: disabled={[e.name for e in extensions if not e.enabled]}')
enabled = [e.name for e in extensions if e.enabled]
disabled = [e.name for e in extensions if not e.enabled]
log.debug(f'Extensions: enabled={enabled} disabled={disabled}')
modules/extra_networks.py
+15 -20
View File
@@ -121,15 +121,6 @@ def activate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str,
p.network_data = extra_network_data
def activate_filtered(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, step=0):
"""activate with text encoder components gated on lora_apply_te; must run before prompt encode so te networks affect embeds"""
apply_te = getattr(p, 'lora_apply_te', None)
if apply_te is None:
apply_te = shared.opts.lora_apply_te
exclude = [] if apply_te else ['text_encoder', 'text_encoder_2', 'text_encoder_3']
activate(p, extra_network_data, step=step, exclude=exclude)
def deactivate(p: StableDiffusionProcessing, extra_network_data: defaultdict[str, list[ExtraNetworkParams]] | None = None, force: bool | None = None):
"""call deactivate for extra networks in extra_network_data in specified order, then call deactivate for all remaining registered networks"""
if p.disable_extra_networks:
@@ -173,20 +164,24 @@ def parse_prompt(prompt: str | None) -> tuple[str, defaultdict[str, list[ExtraNe
return ""
updated_prompt = re.sub(re_extra_net, found, prompt)
updated_prompt = updated_prompt.strip(', ')
return updated_prompt, res
def parse_prompts(prompts: list[str], extra_data: defaultdict[str, list[ExtraNetworkParams]] | None = None):
updated_prompt_list: list[str] = []
extra_data = extra_data or defaultdict(list)
def parse_prompts(
prompts: list[str],
extra_data: defaultdict[str, list[ExtraNetworkParams]] | None = None,
):
updated_prompts: list[str] = []
if extra_data is None:
extra_data = defaultdict(list)
for prompt in prompts:
updated_prompt, parsed_extra_data = parse_prompt(prompt)
if not extra_data:
extra_data = parsed_extra_data
elif parsed_extra_data:
extra_data = parsed_extra_data
else:
pass
updated_prompt_list.append(updated_prompt)
if parsed_extra_data:
for key, values in parsed_extra_data.items():
for item in values:
if item not in extra_data[key]:
extra_data[key].append(item)
return updated_prompt_list, extra_data
updated_prompts.append(updated_prompt)
return updated_prompts, extra_data
modules/face/faceid.py
+1 -1
View File
@@ -216,7 +216,7 @@ def face_id(
p.subseeds = p.all_subseeds[n * p.batch_size:(n+1) * p.batch_size]
p.prompts, p.network_data = extra_networks.parse_prompts(p.prompts, p.network_data)
extra_networks.activate_filtered(p, p.network_data)
extra_networks.activate(p, p.network_data)
ip_model_dict.update({
"prompt": p.prompts[0],
"negative_prompt": p.negative_prompts[0],
modules/framepack/framepack_api.py
-2
View File
@@ -37,7 +37,6 @@ class ReqFramepack(BaseModel):
mp4_opt: str | None = Field(default="crf=16", title="Options", description="Options for the video codec")
mp4_ext: str | None = Field(default="mp4", title="Format", description="Format for the video")
mp4_interpolate: int | None = Field(default=0, title="Interpolation", description="Interpolation for the video")
attention: str | None = Field(default="Default", title="Attention", description="Attention type for the model")
vae_type: str | None = Field(default="Local", title="VAE", description="VAE type for the model")
vlm_enhance: bool | None = Field(default=False, title="VLM enhance", description="Enable VLM enhance")
vlm_model: str | None = Field(default=None, title="VLM model", description="VLM model to use")
@@ -114,7 +113,6 @@ def framepack_post(request: ReqFramepack):
mp4_opt=request.mp4_opt,
mp4_ext=request.mp4_ext,
mp4_interpolate=request.mp4_interpolate,
attention=request.attention,
vae_type=request.vae_type,
vlm_enhance=request.vlm_enhance,
vlm_model=request.vlm_model,
modules/framepack/framepack_install.py
+1 -10
View File
@@ -16,17 +16,8 @@ def rename(src:str, dst:str):
raise e
def install_requirements(attention:str='SDPA'):
def install_requirements():
install('av')
if attention == 'Xformers':
log.debug('FramePack install: xformers')
install('xformers')
elif attention == 'FlashAttention':
log.debug('FramePack install: flash-attn')
install('flash-attn')
elif attention == 'SageAttention':
log.debug('FramePack install: sageattention')
install('sageattention')
def git_clone(git_repo:str, git_dir:str, tmp_dir:str):
modules/framepack/framepack_ui.py
+2 -3
View File
@@ -63,7 +63,6 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
optimized_prompt = gr.Checkbox(label='FP use optimized system prompt', value=True)
use_cfgzero = gr.Checkbox(label='FP enable CFGZero', value=False)
use_preview = gr.Checkbox(label='FP enable Preview', value=True)
attention = gr.Dropdown(label="FP attention", choices=['Default', 'Xformers', 'FlashAttention', 'SageAttention'], value='Default', type='value')
vae_type = gr.Dropdown(label="FP VAE", choices=['Full', 'Tiny', 'Remote'], value='Full', type='value')
with gr.Column(elem_id='framepack-output-column', scale=2) as _column_output:
@@ -88,7 +87,7 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
duration.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
mp4_fps.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
mp4_interpolate.change(fn=change_sections, inputs=[duration, mp4_fps, mp4_interpolate, latent_ws, variant], outputs=[section_html, section_prompt])
btn_load.click(fn=load_model, inputs=[variant, attention], outputs=framepack_outputs)
btn_load.click(fn=load_model, inputs=[variant], outputs=framepack_outputs)
btn_unload.click(fn=unload_model, outputs=framepack_outputs)
receipe_get.click(fn=framepack_load.get_model, inputs=[], outputs=receipe)
receipe_set.click(fn=framepack_load.set_model, inputs=[receipe], outputs=[])
@@ -108,7 +107,7 @@ def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_
use_teacache, use_cfgzero, use_preview,
mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate,
mp4_scale, mp4_upscaler,
attention, vae_type, variant,
vae_type, variant,
vlm_enhance, vlm_model, vlm_system_prompt,
]
modules/framepack/framepack_wrappers.py
+4 -4
View File
@@ -91,11 +91,11 @@ def prepare_prompts(p, init_image, prompt:str, section_prompt:str, num_sections:
return generated_prompts
def load_model(variant, attention):
def load_model(variant):
global loaded_variant # pylint: disable=global-statement
if (shared.sd_model_type != 'hunyuanvideo') or (loaded_variant != variant):
yield gr.update(), gr.update(), 'Verifying FramePack'
framepack_install.install_requirements(attention)
framepack_install.install_requirements()
# framepack_install.git_clone(git_repo=git_repo, git_dir=git_dir, tmp_dir=tmp_dir)
# framepack_install.git_update(git_dir=git_dir, git_commit=git_commit)
# sys.path.append(git_dir)
@@ -114,7 +114,7 @@ def unload_model():
yield gr.update(), gr.update(), 'Model unloaded'
def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_weight, vision_weight, prompt, system_prompt, optimized_prompt, section_prompt, negative_prompt, styles, seed, resolution, duration, latent_ws, steps, cfg_scale, cfg_distilled, cfg_rescale, shift, use_teacache, use_cfgzero, use_preview, mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate, mp4_scale, mp4_upscaler, attention, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt, *_args, **_kwargs):
def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_weight, vision_weight, prompt, system_prompt, optimized_prompt, section_prompt, negative_prompt, styles, seed, resolution, duration, latent_ws, steps, cfg_scale, cfg_distilled, cfg_rescale, shift, use_teacache, use_cfgzero, use_preview, mp4_fps, mp4_codec, mp4_sf, mp4_video, mp4_frames, mp4_thumb, mp4_opt, mp4_ext, mp4_interpolate, mp4_scale, mp4_upscaler, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt, *_args, **_kwargs):
variant = variant or 'Bi-Directional'
if variant == 'None':
log.error('FramePack: no model selected')
@@ -137,7 +137,7 @@ def run_framepack(task_id, _ui_state, init_image, end_image, start_weight, end_w
with call_queue.get_lock():
progress.start_task(task_id)
yield from load_model(variant, attention)
yield from load_model(variant)
if shared.sd_model_type != 'hunyuanvideo':
progress.finish_task(task_id)
yield gr.update(), gr.update(), 'Model load failed'
modules/framepack/pipeline/hunyuan_video_packed.py
+1 -1
View File
@@ -74,7 +74,7 @@ def get_cu_seqlens(text_mask, img_len):
text_len = text_mask.sum(dim=1)
max_len = text_mask.shape[1] + img_len
cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device="cuda")
cu_seqlens = torch.zeros([2 * batch_size + 1], dtype=torch.int32, device=text_mask.device)
for i in range(batch_size):
s = text_len[i] + img_len
modules/generation_parameters_copypaste.py
+17
View File
@@ -296,6 +296,23 @@ def connect_paste(button, local_paste_fields, input_comp, override_settings_comp
res.append(v)
applied[key] = v
else:
if key in ('Sampler', 'Hires sampler') and isinstance(v, str):
from modules import ui_sections
choices, value, _ = ui_sections.sampler_choices(selected=v, same_as_primary=key == 'Hires sampler')
res.append(gr.update(choices=choices, value=value))
applied[key] = v
continue
if getattr(output, 'elem_id', '').endswith('_resize_name') and isinstance(v, str):
from modules import modelloader, ui_sections
modelloader.load_upscalers()
choices = [upscaler.name for upscaler in shared.sd_upscalers]
if output.elem_id.startswith(('control_after', 'control_mask')):
choices = [choice for choice in choices if not choice.lower().startswith('latent')]
if v in choices:
choices, _ = ui_sections.upscaler_choices(choices, selected=v)
res.append(gr.update(choices=choices, value=v))
applied[key] = v
continue
if isinstance(v, str) and v.strip() == '' and key in {'Prompt', 'Negative prompt'}:
debug(f'Paste skip empty: "{key}"')
res.append(gr.update())
modules/gr_hijack.py
+3 -3
View File
@@ -208,13 +208,13 @@ def patch_gradio():
return {"is_generating": False, "data": [], "error": "empty response"}
return response
except GeneratorExit as e:
log.error(f"Gradio queue: events={len(events)} batch={batch} error: {e}")
log.error(f"Gradio queue: events={len(events)} batch={batch} reason=GeneratorExit {e}")
return {"is_generating": False, "data": [None, None, None, None, "cancelled", ""], "error": None}
except Exception as e:
log.error(f"Gradio queue: events={len(events)} batch={batch} error: {e}")
log.error(f"Gradio queue: events={len(events)} batch={batch} reason=Exception {e}")
raise
except BaseException as e:
log.error(f"Gradio queue: events={len(events)} batch={batch} error: {e}")
log.error(f"Gradio queue: events={len(events)} batch={batch} reason=BaseException {e}")
raise
def wrap_blocks_preprocess_data(self, fn_index: int, inputs: list, state: dict):
modules/hashes.py
+5 -3
View File
@@ -1,5 +1,6 @@
import hashlib
import os.path
import threading
from collections import defaultdict
from typing import Literal, TypeAlias, TypedDict
from rich import progress, errors
@@ -29,6 +30,7 @@ cache_filename = os.path.join(data_path, "data", "cache.json")
progress_ok = True
# defaultdict allows for easily using new stores without needing to define them ahead of time
_data: defaultdict[str, HashStore] = defaultdict(HashStore)
cache_lock = threading.Lock()
def load_cache():
@@ -38,9 +40,9 @@ def load_cache():
def save_cache():
# Don't include empty hash stores
filtered = filter(lambda item: len(item[1]) > 0, _data.items())
writefile(dict(filtered), cache_filename)
with cache_lock: # snapshot and write together, so a later save never lands under an earlier snapshot
snapshot = {store: dict(data) for store, data in list(_data.items()) if len(data) > 0} # dict() of a store runs under the GIL, so add_hash from another thread cannot interrupt it
writefile(snapshot, cache_filename)
def cache(store: KnownHashStores | str | None = None) -> HashStore:
modules/history.py
+1 -1
View File
@@ -36,7 +36,7 @@ class Item:
def __str__(self):
if self.latent is not None:
return f'Item(ts="{self.name}" ops={self.ops} latent={self.latent.shape} size={self.size})'
return f'Item(ts="{self.name}" ops={self.ops} latent={list(self.latent.shape)} size={self.size})'
elif self.images is not None:
return f'Item(ts="{self.name}" ops={self.ops} images={len(self.images) if isinstance(self.images, list) else self.images})'
else:
modules/image/grid.py
+5 -4
View File
@@ -80,6 +80,7 @@ def image_grid(imgs: list, batch_size=1, rows: int | None = None, cols: int | No
for i, img in enumerate(params.imgs):
if img is not None:
grid.paste(img, box=(i % params.cols * w, i // params.cols * h))
grid.is_grid = True # flag image as grid
return grid
except Exception as e:
log.error(f'Grid: images={imgs} {e}')
@@ -145,15 +146,15 @@ class GridAnnotation:
def get_font(fontsize: float):
try:
return ImageFont.truetype(shared.opts.font or os.path.join(script_path, "ui", "fonts", "notosans-nerdfont-regular.ttf"), fontsize)
return ImageFont.truetype(shared.opts.font or os.path.join(script_path, "ui", "css", "ubuntu-nerdfont.ttf"), fontsize)
except Exception:
return ImageFont.truetype(os.path.join(script_path, "ui", "fonts", "notosans-nerdfont-regular.ttf"), fontsize)
return ImageFont.truetype(os.path.join(script_path, "ui", "css", "ubuntu-nerdfont.ttf"), fontsize)
def draw_grid_annotations(im: Image.Image, width: int, height: int, x_texts: list[list[GridAnnotation]], y_texts: list[list[GridAnnotation]], margin=0, title: list[GridAnnotation] | None = None):
def wrap(drawing: ImageDraw.ImageDraw, text, font, line_length):
lines = ['']
for word in text.split():
for word in text.split('/\\'):
line = f'{lines[-1]} {word}'.strip()
if drawing.textlength(line, font=font) <= line_length:
lines[-1] = line
@@ -161,7 +162,7 @@ def draw_grid_annotations(im: Image.Image, width: int, height: int, x_texts: lis
lines.append(word)
return lines
def draw_texts(drawing: ImageDraw.ImageDraw, draw_x: float, draw_y: float, lines, initial_fnt: ImageFont.FreeTypeFont, initial_fontsize: int):
def draw_texts(drawing: ImageDraw.ImageDraw, draw_x: float, draw_y: float, lines: list[GridAnnotation], initial_fnt: ImageFont.FreeTypeFont, initial_fontsize: int):
for line in lines:
font = initial_fnt
fontsize = initial_fontsize
modules/image/sharpfin.py
+13 -11
View File
@@ -7,6 +7,7 @@ and Triton GPU acceleration when available.
Non-CUDA devices fall back to PIL/torch.nn.functional automatically.
"""
import os
import sys
import torch
from PIL import Image
@@ -17,6 +18,7 @@ from modules.image.convert import to_tensor, to_pil
_sharpfin_checked = False
_sharpfin_ok = False
_triton_ok = False
debug = log.trace if os.environ.get('SD_PROCESS_DEBUG', None) is not None else lambda *args, **kwargs: None
def check_sharpfin():
@@ -104,7 +106,7 @@ def _scale_pil(scale_fn, tensor, out_res, rk, dev, dt, do_linear, src_h, src_w,
return scale_fn(tensor, out_res, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=True)
except Exception:
_triton_ok = False
log.info("Sharpfin: Triton sparse disabled, using dense path")
log.debug("Sharpfin: Triton sparse disabled, using dense path")
return scale_fn(tensor, out_res, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
# Mixed axis: split into two single-axis resizes
if h > src_h: # H up, W down
@@ -115,7 +117,7 @@ def _scale_pil(scale_fn, tensor, out_res, rk, dev, dt, do_linear, src_h, src_w,
return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=True)
except Exception:
_triton_ok = False
log.info("Sharpfin: Triton sparse disabled, using dense path")
log.debug("Sharpfin: Triton sparse disabled, using dense path")
return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
# H down, W up
use_sparse = _want_sparse(dev, rk, True)
@@ -125,7 +127,7 @@ def _scale_pil(scale_fn, tensor, out_res, rk, dev, dt, do_linear, src_h, src_w,
return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
except Exception:
_triton_ok = False
log.info("Sharpfin: Triton sparse disabled, using dense path")
log.debug("Sharpfin: Triton sparse disabled, using dense path")
intermediate = scale_fn(tensor, (h, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
return scale_fn(intermediate, (h, w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=do_linear, use_sparse=False)
@@ -137,24 +139,24 @@ def resize_pil(image: Image.Image, target_size: tuple[int, int], *, kernel=None,
is_mask = image.mode == 'L'
if (image.width == w) and (image.height == h):
log.debug(f'Resize image: skip={w}x{h} fn={fn}')
# log.debug(f'Resize image: skip={w}x{h} fn={fn}')
return image
from modules import devices
dev = device if device is not None else devices.device
if not allow_sharpfin(dev):
log.debug(f'Resize image: method=PIL source={image.width}x{image.height} target={w}x{h} device={dev} fn={fn}')
debug(f'Resize image: method=PIL source={image.width}x{image.height} target={w}x{h} device={dev} fn={fn}')
return image.resize((w, h), resample=Image.Resampling.LANCZOS)
rk = get_kernel(kernel)
if rk is None:
log.debug(f'Resize image: method=PIL source={image.width}x{image.height} target={w}x{h} kernel=None fn={fn}')
debug(f'Resize image: method=PIL source={image.width}x{image.height} target={w}x{h} kernel=None fn={fn}')
return image.resize((w, h), resample=Image.Resampling.LANCZOS)
from modules.sharpfin.functional import scale
dt = dtype or torch.float16
do_linear = get_linearize(linearize, is_mask=is_mask)
log.debug(f'Resize image: method=sharpfin source={image.width}x{image.height} target={w}x{h} kernel={rk} device={dev} linearize={do_linear} fn={fn}')
debug(f'Resize image: method=sharpfin source={image.width}x{image.height} target={w}x{h} kernel={rk} device={dev} linearize={do_linear} fn={fn}')
tensor = to_tensor(image)
if tensor.dim() == 3:
tensor = tensor.unsqueeze(0)
@@ -182,14 +184,14 @@ def resize_tensor(tensor: torch.Tensor, target_size: tuple[int, int], *, kernel=
dev = devices.device
if not allow_sharpfin(dev):
mode = 'bilinear' if (target_size[0] * target_size[1]) > (tensor.shape[-2] * tensor.shape[-1]) else 'area'
log.debug(f'Resize tensor: method=torch mode={mode} shape={tensor.shape} target={target_size} fn={fn}')
debug(f'Resize tensor: method=torch mode={mode} shape={tensor.shape} target={target_size} fn={fn}')
inp = tensor if tensor.dim() == 4 else tensor.unsqueeze(0)
result = torch.nn.functional.interpolate(inp, size=target_size, mode=mode, antialias=mode != 'area')
return result.squeeze(0) if tensor.dim() == 3 else result
rk = get_kernel(kernel)
if rk is None:
mode = 'bilinear' if (target_size[0] * target_size[1]) > (tensor.shape[-2] * tensor.shape[-1]) else 'area'
log.debug(f'Resize tensor: method=torch mode={mode} shape={tensor.shape} target={target_size} kernel=None fn={fn}')
debug(f'Resize tensor: method=torch mode={mode} shape={tensor.shape} target={target_size} kernel=None fn={fn}')
inp = tensor if tensor.dim() == 4 else tensor.unsqueeze(0)
result = torch.nn.functional.interpolate(inp, size=target_size, mode=mode, antialias=mode != 'area')
return result.squeeze(0) if tensor.dim() == 3 else result
@@ -206,10 +208,10 @@ def resize_tensor(tensor: torch.Tensor, target_size: tuple[int, int], *, kernel=
both_up = (th >= src_h and tw >= src_w)
if both_down or both_up:
use_sparse = _triton_ok and dev.type == 'cuda' and rk.value == 'magic_kernel_sharp_2021' and both_down
log.debug(f'Resize tensor: method=sharpfin shape={tensor.shape} target={target_size} direction={both_up}:{both_down} kernel={rk} sparse={use_sparse} fn={fn}')
debug(f'Resize tensor: method=sharpfin shape={tensor.shape} target={target_size} direction={both_up}:{both_down} kernel={rk} sparse={use_sparse} fn={fn}')
result = scale(tensor, target_size, resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=use_sparse)
else:
log.debug(f'Resize tensor: method=sharpfin shape={tensor.shape} target={target_size} direction={both_up}:{both_down} kernel={rk} sparse=False fn={fn}')
debug(f'Resize tensor: method=sharpfin shape={tensor.shape} target={target_size} direction={both_up}:{both_down} kernel={rk} sparse=False fn={fn}')
intermediate = scale(tensor, (th, src_w), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
result = scale(intermediate, (th, tw), resize_kernel=rk, device=dev, dtype=dt, do_srgb_conversion=linearize, use_sparse=False)
if squeezed:
modules/img2img.py
+12 -13
View File
@@ -163,8 +163,7 @@ def img2img(id_task: str, state: str, mode: int,
vae_type, tiling, hidiffusion,
detailer_enabled, detailer_prompt, detailer_negative, detailer_steps, detailer_strength, detailer_resolution, detailer_classes,
n_iter, batch_size,
guidance_name, guidance_scale, guidance_rescale, guidance_start, guidance_stop,
cfg_scale, cfg_image, cfg_rescale, cfg_true, cfg_adaptive, cfg_end,
cfg_name, cfg_scale, cfg_image, cfg_rescale, cfg_start, cfg_stop, cfg_true, cfg_adaptive,
refiner_start,
clip_skip,
denoising_strength,
@@ -194,6 +193,8 @@ def img2img(id_task: str, state: str, mode: int,
if sampler_index is None:
log.warning('Sampler: invalid')
sampler_index = 0
if hr_sampler_index is None or hr_sampler_index == 'Same as primary':
hr_sampler_index = sampler_index
mode = int(mode)
image = None
@@ -261,13 +262,6 @@ def img2img(id_task: str, state: str, mode: int,
batch_size=batch_size,
n_iter=n_iter,
steps=steps,
guidance_name=guidance_name,
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
guidance_start=guidance_start,
guidance_stop=guidance_stop,
cfg_scale=cfg_scale,
cfg_end=cfg_end,
clip_skip=clip_skip,
width=width,
height=height,
@@ -289,10 +283,6 @@ def img2img(id_task: str, state: str, mode: int,
resize_context=resize_context,
scale_by=scale_by,
denoising_strength=denoising_strength,
cfg_image=cfg_image,
cfg_rescale=cfg_rescale,
cfg_true=cfg_true,
cfg_adaptive=cfg_adaptive,
refiner_start=refiner_start,
inpaint_full_res=inpaint_full_res != 0,
inpaint_full_res_padding=inpaint_full_res_padding,
@@ -306,6 +296,15 @@ def img2img(id_task: str, state: str, mode: int,
grading_shadows_tint=grading_shadows_tint, grading_highlights_tint=grading_highlights_tint, grading_split_tone_balance=grading_split_tone_balance,
grading_vignette=grading_vignette, grading_grain=grading_grain,
grading_lut_file=grading_lut_file.name if grading_lut_file is not None else '', grading_lut_strength=grading_lut_strength,
# guidance
cfg_name=cfg_name,
cfg_scale=cfg_scale,
cfg_image=cfg_image,
cfg_rescale=cfg_rescale,
cfg_start=cfg_start,
cfg_stop=cfg_stop,
cfg_true=cfg_true,
cfg_adaptive=cfg_adaptive,
# refiner
enable_hr=enable_hr,
hr_denoising_strength=hr_denoising_strength,
modules/intel/openvino/__init__.py
+5 -5
View File
@@ -251,14 +251,14 @@ def openvino_compile(gm: GraphModule, *example_inputs, model_hash_str: str | Non
hints = {}
if shared.opts.openvino_accuracy == "performance":
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.PERFORMANCE
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.PERFORMANCE # pylint: disable=c-extension-no-member
elif shared.opts.openvino_accuracy == "accuracy":
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.ACCURACY
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.ACCURACY # pylint: disable=c-extension-no-member
if model_hash_str is not None:
hints['CACHE_DIR'] = shared.opts.openvino_cache_path + '/blob'
core.set_property(hints)
log.debug(f'OpenVINO compile: device={device} backend={shared.opts.cuda_compile_backend} hints={hints} file="{file_name}"')
log.debug(f'OpenVINO compile cache: device={device} backend={shared.opts.cuda_compile_backend} accuracy={shared.opts.openvino_accuracy} hints={hints} hash={model_hash_str} file="{file_name}"')
compiled_model = core.compile_model(om, device)
return compiled_model
@@ -274,9 +274,9 @@ def openvino_compile_cached_model(cached_model_path, *example_inputs):
hints = {'CACHE_DIR': shared.opts.openvino_cache_path + '/blob'}
if shared.opts.openvino_accuracy == "performance":
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.PERFORMANCE
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.PERFORMANCE # pylint: disable=c-extension-no-member
elif shared.opts.openvino_accuracy == "accuracy":
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.ACCURACY
hints[ov_hints.execution_mode] = ov_hints.ExecutionMode.ACCURACY # pylint: disable=c-extension-no-member
core.set_property(hints)
device = get_device()
modules/json_helpers.py
+100 -94
View File
@@ -1,13 +1,52 @@
import os
import sys
import contextlib
import threading
import time
import json
from typing import overload, Literal
import fasteners
import orjson
from modules.logger import log
locking_available = True # used by file read/write locking
path_locks: dict[str, threading.RLock] = {}
path_locks_guard = threading.Lock()
def path_lock(filename: str | os.PathLike[str]) -> threading.RLock:
"""One lock per file path; threads of this process serialize on it, other processes are covered by atomic replace."""
key = os.path.normcase(os.path.realpath(filename))
with path_locks_guard:
lock = path_locks.get(key)
if lock is None:
lock = path_locks[key] = threading.RLock()
with contextlib.suppress(OSError):
os.remove(f"{key}.lock") # left behind by the file lock this replaces
return lock
def read_bytes(filename: str | os.PathLike[str], attempts: int = 5, delay: float = 0.01) -> bytes:
"""Read a whole file, waiting out an open that Windows refuses while a replace of the same name is in flight."""
for attempt in range(attempts):
try:
with open(filename, "rb") as file:
return file.read()
except PermissionError:
if attempt == attempts - 1:
raise
time.sleep(delay)
return b""
def replace_file(source: str, target: str, attempts: int = 10, delay: float = 0.05):
"""os.replace that waits out a target another handle holds open, which Windows reports as a permission error."""
for attempt in range(attempts):
try:
os.replace(source, target)
return
except PermissionError:
if attempt == attempts - 1:
raise
time.sleep(delay)
@overload
@@ -17,51 +56,28 @@ def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool
@overload
def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool = False) -> dict | list: ...
def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool = False, *, as_type="") -> dict | list:
global locking_available # pylint: disable=global-statement
"""Read a JSON file; lock=True serializes with writers of the same path in this process."""
data = {} if as_type == "dict" else []
lock_file = None
locked = False
if lock and locking_available:
with path_lock(filename) if lock else contextlib.nullcontext():
try:
lock_file = fasteners.InterProcessReaderWriterLock(f"{filename}.lock")
lock_file.logger.disabled = True # type: ignore - False positive. Bad typing in Fasteners.
locked = lock_file.acquire_read_lock(blocking=True, timeout=3)
except Exception as err:
lock_file = None
locking_available = False
log.error(f'File read lock: file="{filename}" {err}')
locked = False
try:
# if not os.path.exists(filename):
# return {}
t0 = time.time()
with open(filename, "rb") as file:
b = file.read()
t0 = time.time()
b = read_bytes(filename)
if len(b) == 0:
if not silent:
log.warning(f'Read: file="{filename}" empty')
return {} if as_type == "dict" else []
data = orjson.loads(b) # pylint: disable=no-member
# if type(data) is str:
# data = json.loads(data)
t1 = time.time()
if not silent:
fn = f"{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}" # pylint: disable=protected-access
log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1 - t0:.3f} fn={fn}')
except FileNotFoundError as err:
if not silent:
log.debug(f'Read failed: file="{filename}" {err}')
except Exception as err:
if not silent:
log.error(f'Read failed: file="{filename}" {err}')
try:
if locking_available and lock_file is not None:
lock_file.release_read_lock()
if locked and os.path.exists(f"{filename}.lock"):
os.remove(f"{filename}.lock")
except Exception:
locking_available = False
t1 = time.time()
if not silent:
fn = f"{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}" # pylint: disable=protected-access
log.debug(f'Read: file="{filename}" json={len(data)} bytes={os.path.getsize(filename)} time={t1 - t0:.3f} fn={fn}')
except FileNotFoundError as err:
if not silent:
log.debug(f'Read failed: file="{filename}" {err}')
except Exception as err:
if not silent:
log.error(f'Read failed: file="{filename}" {err}')
if isinstance(data, list) and as_type == "dict":
if not data:
@@ -79,65 +95,55 @@ def readfile(filename: str | os.PathLike[str], silent: bool = False, lock: bool
return data
def writefile(obj: dict | list, filename: str | os.PathLike[str], mode="w", silent=False, atomic=False):
def writefile(obj: dict | list, filename: str | os.PathLike[str], mode="w", silent=False, atomic=True):
"""Write obj as JSON through a temp file and replace; writes to the same path from this process run one at a time, in call order."""
import copy
import tempfile
global locking_available # pylint: disable=global-statement
lock_file = None
locked = False
def default(obj):
log.error(f'Save: file="{filename}" not a valid object: {obj}')
return str(obj)
try:
t0 = time.time()
data = copy.deepcopy(obj) # ensure keys/items aren't added/deleted during json.dumps
for k, v in obj.items() if isinstance(obj, dict) else []: # validate each key-by-key to avoid global exceptions
try:
_tmp = json.dumps(v, indent=2, default=default, allow_nan=False, ensure_ascii=False)
except Exception as err:
if not silent:
log.error(f'Save: file="{filename}" key="{k}" value="{v}" {err}')
del data[k]
output = json.dumps(data, indent=2, default=default)
except Exception as err:
log.error(f'Save failed: file="{filename}" {err}')
return
if mode != "w":
atomic = False # append cannot go through a temp file
try:
if locking_available:
lock_file = fasteners.InterProcessReaderWriterLock(f"{filename}.lock") if locking_available else None
lock_file.logger.disabled = True # type: ignore - False positive. Bad typing in Fasteners.
locked = lock_file.acquire_write_lock(blocking=True, timeout=3) if lock_file is not None else False
except Exception as err:
locking_available = False
lock_file = None
log.error(f'File write lock: file="{filename}" {err}')
locked = False
with path_lock(filename):
try:
t0 = time.time()
snapshot = obj.copy() if isinstance(obj, (dict, list)) else obj # dict.copy and list.copy run under the GIL, so a concurrent insert cannot interrupt them
data = copy.deepcopy(snapshot)
for k, v in list(data.items()) if isinstance(data, dict) else []: # validate each key-by-key to avoid global exceptions
try:
_tmp = json.dumps(v, indent=2, default=default, allow_nan=False, ensure_ascii=False)
except Exception as err:
if not silent:
log.error(f'Save: file="{filename}" key="{k}" value="{v}" {err}')
del data[k]
output = json.dumps(data, indent=2, default=default)
except Exception as err:
log.error(f'Save failed: file="{filename}" {err}')
return
try:
if atomic:
with tempfile.NamedTemporaryFile(mode=mode, encoding="utf8", delete=False, dir=os.path.dirname(filename)) as f:
f.write(output)
f.flush()
os.fsync(f.fileno())
os.replace(f.name, filename)
else:
with open(filename, mode=mode, encoding="utf8") as file:
file.write(output)
t1 = time.time()
if not silent:
datalength = len(data)
log.debug(f'Save: file="{filename}" json={datalength} bytes={len(output)} time={t1 - t0:.3f}')
except Exception as err:
log.error(f'Save failed: file="{filename}" {err}')
try:
if locking_available and lock_file is not None:
lock_file.release_write_lock()
if locked and os.path.exists(f"{filename}.lock"):
os.remove(f"{filename}.lock")
except Exception:
locking_available = False
try:
if atomic:
target = os.path.realpath(filename) # replace the file a symlink points at, not the symlink
fd, temp_name = tempfile.mkstemp(dir=os.path.dirname(target), prefix=f"{os.path.basename(target)}.", suffix=".tmp")
try:
with os.fdopen(fd, mode, encoding="utf8") as f:
f.write(output)
f.flush()
os.fsync(f.fileno())
replace_file(temp_name, target)
except BaseException:
with contextlib.suppress(OSError):
os.remove(temp_name)
raise
else:
with open(filename, mode=mode, encoding="utf8") as file:
file.write(output)
t1 = time.time()
if not silent:
datalength = len(data)
log.debug(f'Save: file="{filename}" json={datalength} bytes={len(output)} time={t1 - t0:.3f}')
except Exception as err:
log.error(f'Save failed: file="{filename}" {err}')
modules/loader.py
+8 -10
View File
@@ -34,7 +34,7 @@ try:
import numpy.random # pylint: disable=W0611,C0411 # this causes failure if numpy version changed
def obj2sctype(obj):
return np.dtype(obj).type
if np.__version__.startswith('2.'): # monkeypatch for np==1.2 compatibility
if str(np.__version__).startswith('2.'): # monkeypatch for np==1.2 compatibility
np.obj2sctype = obj2sctype # noqa: NPY201
np.bool8 = np.bool
np.float_ = np.float64 # noqa: NPY201
@@ -133,13 +133,6 @@ if ".dev" in torch.__version__ or "+git" in torch.__version__:
torch.__version__ = re.search(r'[\d.]+[\d]', torch.__version__).group(0)
timer.startup.record("torch")
try:
from modules.sd_hijack_triton import install as install_autotune_report # pylint: disable=ungrouped-imports
install_autotune_report()
except Exception as e:
log.warning(f'Triton logging: {e}')
timer.startup.record("triton")
try:
import bitsandbytes # pylint: disable=unused-import
_bnb = True
@@ -201,8 +194,6 @@ except Exception as e:
_onnx = False
timer.startup.record("onnx")
timer.startup.record("fastapi")
import gradio # pylint: disable=W0611,C0411
timer.startup.record("gradio")
errors.install([gradio])
@@ -237,6 +228,13 @@ diffusers.utils.import_utils._sdnq_available = True # pylint: disable=protected-
diffusers.utils.import_utils._sdnq_version = sdnq.__version__ # pylint: disable=protected-access
timer.startup.record("sdnq")
try:
from modules.sd_hijack_triton import install as install_autotune_report # pylint: disable=ungrouped-imports
install_autotune_report()
except Exception as e:
log.warning(f'Triton logging: {e}')
timer.startup.record("triton")
try:
import pillow_jxl # pylint: disable=W0611,C0411
except Exception:
modules/logger.py
+30 -15
View File
@@ -182,21 +182,14 @@ def setup_logging(debug=None, trace=None, filename=None):
render_options = render_options.update_height(height=render_options.height - self.top - self.bottom)
lines = console.render_lines(self.renderable, render_options, style=style, pad=False)
_Segment = Segment
left = _Segment(" " * self.left, style) if self.left else None
right = [_Segment.line()]
blank_line: list[Segment] | None = None
if self.top:
blank_line = [_Segment(f'{" " * width}\\n', style)]
yield from blank_line * self.top
if left:
for line in lines:
yield left
yield from line
yield from right
else:
for line in lines:
yield from line
yield from right
for line in lines: # self.left is forced to 0 above, so no left-padding segment is ever emitted
yield from line
yield from right
if self.bottom:
blank_line = blank_line or [_Segment(f'{" " * width}\\n', style)]
yield from blank_line * self.bottom
@@ -232,13 +225,26 @@ def setup_logging(debug=None, trace=None, filename=None):
log.setLevel(logging.DEBUG) # log to file is always at level debug for facility `sd`
log.print = rprint
# use only the 16 standard ANSI color names (+ dim/bold modifiers) so the theme renders correctly on basic 16-color terminals too
theme = Theme({
"traceback.border": "black",
"inspect.value.border": "black",
"traceback.border.syntax_error": "dark_red",
"logging.level.info": "blue_violet",
"logging.level.debug": "purple4",
"logging.level.trace": "dark_blue",
"traceback.border.syntax_error": "red",
"logging.level.trace": "dim cyan",
"logging.level.debug": "cyan",
"logging.level.info": "bright_cyan",
"logging.level.warning": "yellow",
"logging.level.error": "red",
"logging.level.critical": "bold bright_red",
"repr.attrib_name": "bright_white",
"repr.attrib_value": "cyan",
"repr.str": "bright_cyan",
"repr.none": "yellow",
"repr.number": "bright_yellow",
"repr.bool_true": "green",
"repr.bool_false": "bright_red",
})
Padding.__rich_console__ = override_padding
@@ -263,7 +269,7 @@ def setup_logging(debug=None, trace=None, filename=None):
log_filter = LogFilter()
# handlers
rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=False, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=level, console=console)
rh = RichHandler(show_time=True, omit_repeated_times=False, show_level=True, show_path=False, markup=True, rich_tracebacks=True, log_time_format='%H:%M:%S-%f', level=level, console=console)
if trace:
rh.formatter = logging.Formatter('[%(module)s][%(pathname)s:%(lineno)d] %(message)s')
rh.addFilter(log_filter)
@@ -312,3 +318,12 @@ def setup_logging(debug=None, trace=None, filename=None):
logging.getLogger("torch").setLevel(logging.DEBUG)
else:
logging.getLogger("torch").setLevel(logging.WARNING)
if __name__ == "__main__":
setup_logging(debug=True, trace=False, filename=None)
for l in [logging.TRACE, logging.DEBUG, logging.INFO, logging.WARNING, logging.ERROR, logging.CRITICAL]:
log.log(l, f"Test log level: {logging.getLevelName(l)}")
values = [None, True, False, "yes", "no", "sd.next", 1, 0, 1.0, [1,2,3], {"key": "value"}, (1,2), {1,2}, object()]
for v in values:
log.info(f"Test values: {type(v).__name__}={v}")
modules/lora/extra_networks_lora.py
+33 -20
View File
@@ -27,7 +27,7 @@ def get_stepwise(param, step, steps): # from https://github.com/cheald/sd-webui-
if m[1][-1] <= 1.0:
step = step / (max_steps - step_offset) if max_steps > 0 else 1.0
v = np.interp(step, m[1], m[0])
debug_log(f"Network load: type=LoRA step={step} steps={max_steps} v={v}")
debug_log(f"LoRA: stepwise step={step} steps={max_steps} v={v}")
return v
else:
return m
@@ -54,7 +54,7 @@ def prompt(p):
all_tags = list(set(all_tags))
all_tags = [t for t in all_tags if t not in p.prompt]
if len(all_tags) > 0:
log.debug(f"Network load: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply")
log.debug(f"Network tags: type=LoRA tags={all_tags} max={shared.opts.lora_apply_tags} apply")
all_tags = ', '.join(all_tags)
p.extra_generation_params["LoRA tags"] = all_tags
if '_tags_' in p.prompt:
@@ -98,6 +98,7 @@ def parse(p, params_list, step=0):
unet_multipliers = []
dyn_dims = []
lora_modules = []
block_specs = []
for params in params_list:
name = params.positional[0]
@@ -131,6 +132,7 @@ def parse(p, params_list, step=0):
te_multipliers.append(te_multiplier)
unet_multipliers.append(unet_multiplier)
dyn_dims.append(dyn_dim)
block_specs.append(params.named.get('lbw', None)) # per-block strength; resolved per layer by lora_blocks
lora_module = []
name_lower = params.positional[0].lower()
@@ -150,7 +152,7 @@ def parse(p, params_list, step=0):
lora_modules.append(lora_module)
return names, te_multipliers, unet_multipliers, dyn_dims, lora_modules
return names, te_multipliers, unet_multipliers, dyn_dims, lora_modules, block_specs
def unload_diffusers():
@@ -161,7 +163,7 @@ def unload_diffusers():
pass
if hasattr(shared.sd_model, "unload_lora_weights"):
try:
shared.sd_model.unload_lora_weights() # fails for non-CLIP models
shared.sd_model.unload_lora_weights()
except Exception:
pass
@@ -174,12 +176,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
self.model = None
self.errors = {}
def signature(self, names: list[str], te_multipliers: list, unet_multipliers: list):
return [f'{name}:{te}:{unet}' for name, te, unet in zip(names, te_multipliers, unet_multipliers, strict=False)]
def signature(self, names: list[str], te_multipliers: list, unet_multipliers: list, block_specs: list | None = None):
specs = block_specs if block_specs else [None] * len(names)
return [f'{name}:{te}:{unet}' + (f':lbw={str(spec).strip().lower()}' if spec else '') for name, te, unet, spec in zip(names, te_multipliers, unet_multipliers, specs, strict=False)]
def changed(self, requested: list[str], include: list[str] | None = None, exclude: list[str] | None = None) -> tuple[bool, str]:
from modules.lora import lora_sdnq, lora_stack
requested = requested + [f'stack={lora_stack.signature()}{lora_sdnq.signature()}'] # settings-only stack or mechanism changes must re-trigger activation
if shared.opts.lora_force_reload:
debug_log(f'Network check: type=LoRA requested={requested} status="forced"')
debug_log(f'LoRA check requested={requested} status="forced"')
return True, "forced"
sd_model = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
if sd_model is None:
@@ -195,15 +200,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
if len(requested) != len(loaded):
sd_model.loaded_loras.clear() # single-entry cache: any activation invalidates state recorded under other filter keys
sd_model.loaded_loras[key] = requested
debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="num changed"')
debug_log(f'LoRA check key="{key}" requested={requested} loaded={loaded} status="num changed"')
return True, "num changed"
for req, load in zip(requested, loaded, strict=False):
if req != load:
sd_model.loaded_loras.clear()
sd_model.loaded_loras[key] = requested
debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="content changed"')
debug_log(f'LoRA check key="{key}" requested={requested} loaded={loaded} status="content changed"')
return True, "content changed"
debug_log(f'Network check: type=LoRA key="{key}" requested={requested} loaded={loaded} status="same"')
debug_log(f'LoRA check key="{key}" requested={requested} loaded={loaded} status="same"')
return False, "none"
def activate(self, p, params_list, step=0, include=None, exclude=None): # pylint: disable=arguments-differ
@@ -218,15 +223,20 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
if len(params_list) > 0 and not self.active: # activate patches once
self.active = True
self.model = shared.opts.sd_model_checkpoint
names, te_multipliers, unet_multipliers, dyn_dims, lora_modules = parse(p, params_list, step)
requested = self.signature(names, te_multipliers, unet_multipliers)
names, te_multipliers, unet_multipliers, dyn_dims, lora_modules, block_specs = parse(p, params_list, step)
requested = self.signature(names, te_multipliers, unet_multipliers, block_specs)
reason = ''
load_method, load_reason = lora_overrides.get_method()
from modules.lora import lora_stack
if load_method != 'native' and lora_stack.mode() != 'sum':
log.warning(f'Network stack: mode={lora_stack.mode()} method={load_method} fallback=sum')
if load_method != 'native' and any(block_specs):
log.warning(f'Network blocks: method={load_method} fallback=none')
if debug:
import sys
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
debug_log(f'Network load: type=LoRA include={include} exclude={exclude} method={load_method} reason="{load_reason}" requested={requested} fn={fn}')
debug_log(f'LoRA load: include={include} exclude={exclude} method={load_method} reason="{load_reason}" requested={requested} fn={fn}')
if load_method == 'diffusers':
has_changed, reason = self.changed(requested)
@@ -238,7 +248,7 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
if hasattr(sd_model, 'disable_lora'):
try:
sd_model.disable_lora()
log.info('Network unload: type=LoRA mode=diffusers')
log.info('Network unload: type=LoRA method=diffusers disable')
except Exception as e:
log.error(f'Network unload: type=LoRA {e}')
sd_models.set_diffuser_offload(shared.sd_model, op="model")
@@ -249,15 +259,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
has_changed = lora_nunchaku.load_nunchaku(names, unet_multipliers)
else: # native
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, activate=False) # load only, activation below honors include/exclude
lora_load.network_load(names, te_multipliers, unet_multipliers, dyn_dims, block_specs=block_specs, activate=False) # load only, activation below honors include/exclude
has_changed, reason = self.changed(requested, include, exclude)
if has_changed:
jobid = shared.state.begin('LoRA')
if len(l.previously_loaded_networks) > 0:
log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={"fuse" if shared.opts.lora_fuse_native else "backup"}')
log.info(f'Network unload: type=LoRA networks={[n.name for n in l.previously_loaded_networks]} mode={networks.effective_mode()}')
networks.network_deactivate(include, exclude)
networks.network_activate(include, exclude)
debug_log(f'Network change: type=LoRA previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]}')
debug_log(f'LoRA change: previous={[n.name for n in l.previously_loaded_networks]} current={[n.name for n in l.loaded_networks]}')
if len(include) == 0:
l.previously_loaded_networks = l.loaded_networks.copy()
shared.state.end(jobid)
@@ -265,12 +275,15 @@ class ExtraNetworkLora(extra_networks.ExtraNetwork):
if len(l.loaded_networks) > 0 and (len(networks.applied_layers) > 0 or load_method=='diffusers' or load_method=='nunchaku') and step == 0:
infotext(p)
prompt(p)
if has_changed and len(include) == 0: # print only once
sd_model = shared.sd_model.pipe if hasattr(shared.sd_model, 'pipe') else shared.sd_model
if len(include) == 0: # print only once
actual_method = 'native' if any(len(n.modules) > 0 for n in l.loaded_networks) else load_method
log.info(f'Network load: type=LoRA networks={[n.name for n in l.loaded_networks]} load={load_method}({load_reason}) method={actual_method} mode={"fuse" if shared.opts.lora_fuse_native else "backup"} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} reason="{reason}"')
stack = lora_stack.signature() if actual_method == 'native' else 'sum' # non-native paths always combine as sum
log.info(f'Network status: type=LoRA networks={[n.name for n in l.loaded_networks]} method={actual_method}({load_reason}) mode={networks.effective_mode()} stack={stack} te={te_multipliers} unet={unet_multipliers} time={l.timer.summary} changed={has_changed} reason="{reason}"')
def deactivate(self, p, force=False):
def deactivate(self, p, force=False): # pylint: disable=unused-argument
if len(lora_diffusers.diffuser_loaded) > 0 and (shared.opts.lora_force_reload or force):
log.debug(f'Network unload: type=LoRA method=diffusers loaded={len(lora_diffusers.diffuser_loaded)} opts={shared.opts.lora_force_reload} force={force}')
unload_diffusers()
if force:
networks.network_deactivate()
modules/lora/lora_apply.py
+23 -6
View File
@@ -5,6 +5,7 @@ import time
from typing import TYPE_CHECKING
import torch
from modules.lora import lora_common as l
from modules.lora import lora_stack
from modules import shared, devices, errors
from modules.logger import log
@@ -16,7 +17,7 @@ if TYPE_CHECKING:
re_network_name = re.compile(r"(.*)\s*\([0-9a-fA-F]+\)")
def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple):
def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, wanted_names: tuple, fuse: bool):
backup_size = 0
if len(l.loaded_networks) > 0 and network_layer_name is not None and any([net.modules.get(network_layer_name, None) for net in l.loaded_networks]): # noqa: C419 # pylint: disable=R1729
t0 = time.time()
@@ -24,7 +25,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
weights_backup = getattr(self, "network_weights_backup", None)
bias_backup = getattr(self, "network_bias_backup", None)
if weights_backup is not None or bias_backup is not None:
if (shared.opts.lora_fuse_native and not isinstance(weights_backup, bool)) or (not shared.opts.lora_fuse_native and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly
if (fuse and not isinstance(weights_backup, bool)) or (not fuse and isinstance(weights_backup, bool)): # invalidate so we can change direct/backup on-the-fly
weights_backup = None
bias_backup = None
self.network_weights_backup = weights_backup
@@ -33,7 +34,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
if weights_backup is None and wanted_names != (): # pylint: disable=C1803
weight = getattr(self, 'weight', None)
self.network_weights_backup = None
if shared.opts.lora_fuse_native:
if fuse:
self.network_weights_backup = True
else:
self.network_weights_backup = weight.clone().to(devices.cpu)
@@ -53,7 +54,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
if bias_backup is None:
if getattr(self, 'bias', None) is not None:
if shared.opts.lora_fuse_native:
if fuse:
self.network_bias_backup = True
else:
bias_backup = self.bias.clone()
@@ -67,7 +68,7 @@ def network_backup_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Gr
return backup_size
def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None):
def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, network_layer_name: str, use_previous: bool = False, *, elimit: Callable[[], None] | None = None, per_net: bool = False):
if shared.opts.diffusers_offload_mode == "none":
try:
self.to(devices.device)
@@ -75,6 +76,9 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
pass
batch_updown = None
batch_ex_bias = None
stack_deltas = None
if per_net or (lora_stack.mode() in lora_stack.DENSE_MODES and network_layer_name is not None and not network_layer_name.startswith('lora_te')):
stack_deltas = [] # collect per-net deltas; combined after the loop unless the caller wants them separate (bias deltas stay summed)
loaded = l.loaded_networks if not use_previous else l.previously_loaded_networks
for net in loaded:
module = net.modules.get(network_layer_name, None)
@@ -107,7 +111,9 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
del weight
if updown is not None:
if batch_updown is not None:
if stack_deltas is not None:
stack_deltas.append((net.name, updown.to(devices.device)))
elif batch_updown is not None:
batch_updown += updown.to(batch_updown.device)
else:
batch_updown = updown.to(devices.device)
@@ -136,6 +142,17 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
if elimit is not None:
elimit()
continue
if per_net:
return stack_deltas, batch_ex_bias
if stack_deltas is not None and stack_deltas:
if len(stack_deltas) >= 2:
t0 = time.time()
batch_updown = lora_stack.combine(stack_deltas, network_layer_name)
l.timer.calc += time.time() - t0
else:
batch_updown = stack_deltas[0][1]
if shared.opts.diffusers_offload_mode == "sequential":
batch_updown = batch_updown.to(devices.cpu)
return batch_updown, batch_ex_bias
modules/lora/lora_blocks.py
+348
View File
@@ -0,0 +1,348 @@
"""Per-block LoRA strength: <lora:name:1.0:lbw=VALUE>.
Each targeted layer maps to one slot of a per-architecture weight vector and
the network's multiplier is scaled by that slot. Slot 0 is BASE: on unet
architectures it covers the text encoder and the unet layers outside the
block chain, on transformer architectures the layers outside the block
chain(s). The remaining slots follow the merge block-weight layout on unet
architectures (26 on sd, 20 on sdxl: input blocks, mid, output blocks) and
the transformer chain(s) in depth order elsewhere, with chain lengths
scanned from the live network_layer_mapping rather than hardcoded.
VALUE is a preset name (case-insensitive), a single number broadcast to
every slot, or a comma list with one number per slot. Named presets force
BASE to 1.0, since the merge tables carry 0 there with merge semantics, and
stretch onto the block count of the current model; classic segment names
(INS, OUTALL, ...) generate from ranges, so they also work on transformer
chains via thirds, and DOUBLE/SINGLE mute one chain on two-chain
architectures. Explicit vectors are taken verbatim at the slot count, with
the a1111 17-slot (sd) and 12-slot (sdxl) layouts accepted and expanded,
omitted slots neutral. A value that fits nothing is ignored with a warning
and the network applies at its plain strength.
"""
import re
from modules import shared
from modules.logger import log
from modules.lora import lora_common as l
UNET_ARCHES = ('sd', 'sdxl')
CHAINS = { # arch -> anchored tail prefixes, one per chain, in depth order
'sd3': ('transformer_blocks_',),
'anima': ('transformer_blocks_',),
'f1': ('transformer_blocks_', 'single_transformer_blocks_'),
'f2': ('transformer_blocks_', 'single_transformer_blocks_'),
'chroma': ('transformer_blocks_', 'single_transformer_blocks_'),
'zimage': ('layers_',),
'ernieimage': ('layers_',),
'krea2': ('blocks_',),
}
CLASSIC = ('ALL', 'NONE', 'INALL', 'INS', 'IND', 'MIDD', 'OUTALL', 'OUTD', 'OUTS')
CHAIN_NAMES = ('DOUBLE', 'SINGLE')
SD1_17 = (0, 2, 3, 5, 6, 8, 9, 13, 17, 18, 19, 20, 21, 22, 23, 24, 25) # BASE, IN01, IN02, IN04, IN05, IN07, IN08, MID, OUT03..OUT11
SDXL_12 = (0, 5, 6, 8, 9, 10, 11, 12, 13, 14, 15, 16) # BASE, IN04, IN05, IN07, IN08, MID, OUT00..OUT05
VECTOR_MEMO_CAP = 64
MISS = object()
re_down = re.compile(r'^down_blocks_(\d+)_(resnets|attentions|downsamplers)_(\d+)')
re_up = re.compile(r'^up_blocks_(\d+)_(resnets|attentions|upsamplers)_(\d+)')
re_chain_index = re.compile(r'^(\d+)')
state: dict = {'stamp': None, 'layout': None, 'index': {}, 'vectors': {}}
warned: set = set()
def warn_once(key, message):
if key not in warned:
warned.add(key)
log.warning(message)
def build_unet_layout(arch, mapping):
down, up = -1, -1
for key in mapping:
if not key.startswith('lora_unet_'):
continue
tail = key[len('lora_unet_'):]
m = re_down.match(tail)
if m is not None:
down = max(down, int(m.group(1)))
continue
m = re_up.match(tail)
if m is not None:
up = max(up, int(m.group(1)))
if down < 0 or up < 0:
return None
n_in = 3 * (down + 1) # conv_in plus two pairs and a sampler slot per group: the compvis input_blocks count
n_out = 3 * (up + 1)
n = 2 + n_in + n_out
return {
'arch': arch, 'kind': 'unet', 'n': n, 'n_in': n_in,
'ins': list(range(1, 1 + n_in)),
'mids': [1 + n_in],
'outs': list(range(2 + n_in, n)),
}
def build_dit_layout(arch, mapping):
prefixes = CHAINS.get(arch)
if prefixes is None:
return None
counts = [0 for _ in prefixes]
for key in mapping:
if not key.startswith('lora_transformer_'):
continue
tail = key[len('lora_transformer_'):]
for i, prefix in enumerate(prefixes):
if tail.startswith(prefix):
m = re_chain_index.match(tail[len(prefix):])
if m is not None:
counts[i] = max(counts[i], int(m.group(1)) + 1)
break
total = sum(counts)
if total == 0:
return None
chains = []
offset = 0
for prefix, count in zip(prefixes, counts, strict=False):
chains.append((prefix, count, offset))
offset += count
n = 1 + total
blocks = list(range(1, n))
return {
'arch': arch, 'kind': 'dit', 'n': n, 'chains': chains,
'ins': [s for i, s in enumerate(blocks) if i * 3 // total == 0],
'mids': [s for i, s in enumerate(blocks) if i * 3 // total == 1],
'outs': [s for i, s in enumerate(blocks) if i * 3 // total == 2],
}
def layout():
sd_model = getattr(shared, 'sd_model', None)
mapping = getattr(sd_model, 'network_layer_mapping', None) if sd_model is not None else None
if not mapping:
return None
arch = shared.sd_model_type
stamp = (arch, id(mapping))
if state['stamp'] == stamp:
return state['layout']
state['stamp'] = stamp
state['layout'] = build_unet_layout(arch, mapping) if arch in UNET_ARCHES else build_dit_layout(arch, mapping)
state['index'].clear()
state['vectors'].clear()
return state['layout']
def classify(sd_key, lay):
if sd_key.startswith('lora_te'):
return 0 if lay['kind'] == 'unet' else None # BASE covers the TE on unet arches; transformer vectors do not model the TE
if sd_key.startswith('lora_llm_adapter_'):
return None
if lay['kind'] == 'unet':
if not sd_key.startswith('lora_unet_'):
return None
tail = sd_key[len('lora_unet_'):]
m = re_down.match(tail)
if m is not None:
slot = 1 + 3 * int(m.group(1)) + (2 if m.group(2) == 'downsamplers' else int(m.group(3)))
return 1 + slot
m = re_up.match(tail)
if m is not None:
slot = 3 * int(m.group(1)) + (2 if m.group(2) == 'upsamplers' else int(m.group(3)))
return 2 + lay['n_in'] + slot
if tail.startswith('mid_block'):
return 1 + lay['n_in']
if tail.startswith('conv_in'):
return 1 # IN00
if tail.startswith('conv_out') or tail.startswith('conv_norm_out'):
return lay['n'] - 1 # the compvis out group belongs to the last output block
return 0 # time_embedding, add_embedding and other non-block leaves
if not sd_key.startswith('lora_transformer_'):
return None
tail = sd_key[len('lora_transformer_'):]
for prefix, count, offset in lay['chains']:
if tail.startswith(prefix):
m = re_chain_index.match(tail[len(prefix):])
if m is not None and int(m.group(1)) < count:
return 1 + offset + int(m.group(1))
return 0
return 0 # embedders, projections, refiners and other non-chain layers
def block_index(sd_key):
lay = layout()
if lay is None:
return None
cached = state['index'].get(sd_key, MISS)
if cached is not MISS:
return cached
idx = classify(sd_key, lay)
state['index'][sd_key] = idx
return idx
def fill_band(vec, slots, lo, hi):
k = len(slots)
for i, s in enumerate(slots):
if lo * k <= i < hi * k:
vec[s] = 1.0
def classic_vector(name, lay):
if name == 'ALL':
return [1.0] * lay['n']
vec = [0.0] * lay['n']
if name == 'NONE':
return vec
vec[0] = 1.0
if name == 'INALL':
fill_band(vec, lay['ins'], 0.0, 1.0)
elif name == 'INS': # shallow half of the input side
fill_band(vec, lay['ins'], 0.0, 0.5)
elif name == 'IND': # deep half of the input side
fill_band(vec, lay['ins'], 0.5, 1.0)
elif name == 'MIDD': # the middle of the network: deep input half, mid, deep output half
fill_band(vec, lay['ins'], 0.5, 1.0)
fill_band(vec, lay['mids'], 0.0, 1.0)
fill_band(vec, lay['outs'], 0.0, 0.5)
elif name == 'OUTALL':
fill_band(vec, lay['outs'], 0.0, 1.0)
elif name == 'OUTD': # deep half of the output side, nearest the mid
fill_band(vec, lay['outs'], 0.0, 0.5)
elif name == 'OUTS': # shallow half of the output side, nearest the image
fill_band(vec, lay['outs'], 0.5, 1.0)
return vec
def chain_vector(name, lay):
chains = lay.get('chains') or []
if len(chains) != 2:
return None
vec = [1.0] * lay['n']
keep = 0 if name == 'DOUBLE' else 1
for i, (_prefix, count, offset) in enumerate(chains):
val = 1.0 if i == keep else 0.0
for s in range(1 + offset, 1 + offset + count):
vec[s] = val
return vec
def stretch(src, k):
if k == len(src):
return [float(v) for v in src]
out = []
for i in range(k):
x = i * (len(src) - 1) / (k - 1) if k > 1 else 0.0
lo = int(x)
hi = min(lo + 1, len(src) - 1)
f = x - lo
out.append(float(src[lo]) * (1.0 - f) + float(src[hi]) * f)
return out
def preset_vector(name, lay):
from modules.merging.merge_presets import BLOCK_WEIGHTS_PRESETS, SDXL_BLOCK_WEIGHTS_PRESETS
if name in CHAIN_NAMES:
return chain_vector(name, lay)
if name in CLASSIC:
return classic_vector(name, lay)
if lay['arch'] == 'sdxl':
src = SDXL_BLOCK_WEIGHTS_PRESETS.get(name) or SDXL_BLOCK_WEIGHTS_PRESETS.get('SDXL_' + name)
if src is not None:
return [1.0] + [float(v) for v in src[1:]] # merge tables carry 0 in the BASE slot; a preset must leave the TE alone
if name.startswith('SDXL_'):
return None # explicitly arch-tagged, not reinterpreted elsewhere
src = BLOCK_WEIGHTS_PRESETS.get(name)
if src is None:
return None
if lay['arch'] == 'sd':
return [1.0] + [float(v) for v in src[1:]]
return [1.0] + stretch(src[1:], lay['n'] - 1)
def parse_vector(parts, lay):
try:
vals = [float(x) for x in parts]
except ValueError:
return None
n = lay['n']
if len(vals) == n:
return vals
if len(vals) == n - 1:
return [1.0] + vals
legacy = SD1_17 if lay['arch'] == 'sd' else (SDXL_12 if lay['arch'] == 'sdxl' else None)
if legacy is not None and len(vals) == len(legacy):
vec = [1.0] * n # slots the a1111 layouts omit stay neutral
for slot, v in zip(legacy, vals, strict=False):
vec[slot] = v
return vec
return None
def resolve(spec):
"""Resolve a raw lbw value into a slot vector for the current model, or None when it fits nothing."""
lay = layout()
if lay is None:
return None
raw = str(spec).strip()
key = raw.lower()
if key in state['vectors']:
return state['vectors'][key]
if len(state['vectors']) > VECTOR_MEMO_CAP:
state['vectors'].clear()
vec = None
if ',' in raw:
vec = parse_vector([x.strip() for x in raw.split(',')], lay)
if vec is None:
warn_once(f'lbw-vector:{key}:{lay["arch"]}', f'Network blocks: value="{raw}" arch={lay["arch"]} expected={lay["n"]} fallback=none')
else:
try:
vec = [float(raw)] * lay['n']
except ValueError:
vec = preset_vector(raw.upper(), lay)
if vec is None:
warn_once(f'lbw-name:{key}:{lay["arch"]}', f'Network blocks: preset="{raw}" arch={lay["arch"]} fallback=none')
if vec is not None:
log.info(f'Network blocks: value="{raw}" arch={lay["arch"]} slots={lay["n"]} range={min(vec):.2f}-{max(vec):.2f}')
state['vectors'][key] = vec
return vec
def factor(sd_key, net):
"""Per-layer scale from a network's block vector; 1.0 whenever the vector does not apply."""
try:
spec = getattr(net, 'block_spec', None)
if not spec:
return 1.0
vec = resolve(spec)
if vec is None:
return 1.0
idx = block_index(sd_key)
if idx is None:
return 1.0
return float(vec[idx])
except Exception as e:
warn_once('lbw-error', f'Network blocks: {e} fallback=none')
return 1.0
def net_signature(net):
"""Normalized spec of one network, or None; joins content identities such as the factor cache signature."""
spec = getattr(net, 'block_spec', None)
if not spec:
return None
return str(spec).strip().lower()
def active():
return any(getattr(net, 'block_spec', None) for net in l.loaded_networks)
def signature():
"""Identity suffix for the per-module apply stamp; empty while no loaded network carries block weights."""
specs = [f'{net.name}:{net_signature(net)}' for net in l.loaded_networks if getattr(net, 'block_spec', None)]
if len(specs) == 0:
return ''
return '|lbw=' + ','.join(specs)
modules/lora/lora_calib.py
+226
View File
@@ -0,0 +1,226 @@
"""Per-checkpoint activation calibration for svd hosting on quantized layers.
Plain svd truncation of a hosted delta is optimal in weight space but not in
output space: transformer activations concentrate energy in a few input
channels (per-channel RMS spreads by one to three orders of magnitude), so
the directions that matter most for the output are not the largest in
Frobenius norm. Scaling the delta by per-channel input RMS before the svd
and folding the inverse scale into the down factor spends the same rank
budget on output error instead; measured on real checkpoints this raises
output-delta retention by ~0.05 at rank 256 and ~0.09 at rank 64, most on
MLP down projections whose inputs carry the largest outlier channels.
Statistics come from the model's own forwards: when a sub-8-bit SDNQ model
loads and no calibration is cached for it, streaming sum-of-squares hooks
attach to its quantized linears, accumulate during normal generations,
persist, and go inert. Persist fires when every layer reaches the token
quota, or at a bounded number of denoiser forwards for models where some
projections take pooled or modulation vectors (a few tokens per forward)
and could never reach an absolute quota; layers still under a small token
floor at the deadline are omitted and stay on plain truncation. Cached
statistics load at model load and sit on each layer as ``sdnq_calib_rms``;
the hosting path reads them through ``rms_for``. Capture is skipped when
the model is compiled (hooks would break the graph) and everything is
gated by the ``lora_sdnq_host_calib`` option.
"""
import os
from typing import Optional, TypedDict
import torch
from modules import paths, shared, script_callbacks
from modules.logger import log
class CaptureRecord(TypedDict):
m: torch.nn.Module
ss: Optional[torch.Tensor]
n: int
done: bool
class CaptureState(TypedDict):
model: Optional[str]
recs: dict[str, CaptureRecord]
handles: list[torch.utils.hooks.RemovableHandle]
forwards: int
complete: bool
TOKENS_DONE = 65536
FORWARDS_DEADLINE = 48 # ~2 generations; token-rich layers normally finish their quota well inside it
TOKENS_FLOOR = 32 # below this mass the rms estimate is noise; the layer is omitted and stays on plain truncation
calib_root = os.path.join(paths.models_path, 'calibration')
capture: CaptureState = {'model': None, 'recs': {}, 'handles': [], 'forwards': 0, 'complete': False}
def enabled():
return bool(getattr(shared.opts, 'lora_sdnq_host_calib', False))
def calib_file(model_name):
key = model_name.replace('/', '--').replace('\\', '--').replace(':', '-')
return os.path.join(calib_root, f'{key}.safetensors')
def checkpoint_name(sd_model):
info = getattr(sd_model, 'sd_checkpoint_info', None)
return getattr(info, 'name', None)
def denoiser_root(sd_model):
"""The model's denoiser component, transformer first, unet otherwise."""
root = getattr(sd_model, 'transformer', None)
return root if root is not None else getattr(sd_model, 'unet', None)
def eligible_modules(sd_model):
"""Sub-8-bit 2-D SDNQ linears of the model's denoiser: the layers hosting applies to."""
root = denoiser_root(sd_model)
if root is None:
return []
from sdnq.common import dtype_dict
out = []
for name, m in root.named_modules():
deq = getattr(m, 'sdnq_dequantizer', None)
if deq is None or len(deq.original_shape) != 2:
continue
if dtype_dict[deq.weights_dtype]['num_bits'] >= 8:
continue
out.append((name, m))
return out
def detach_capture():
for h in capture['handles']:
h.remove()
capture['handles'].clear()
capture['recs'].clear()
capture['model'] = None
capture['forwards'] = 0
capture['complete'] = False
def deadline_hook(module, hook_args): # pylint: disable=unused-argument
"""Count denoiser forwards and close capture at the deadline.
Layers taking pooled or modulation vectors see a few tokens per forward
and can never reach the token quota; a global forward count bounds
capture for them and for modules the generation path never runs.
"""
if capture['complete']:
return
capture['forwards'] += 1
if capture['forwards'] >= FORWARDS_DEADLINE:
persist()
def hook_for(rec, in_features):
def hook(module, hook_args): # pylint: disable=unused-argument
if rec['done'] or capture['complete']:
return
x = hook_args[0] if hook_args else None
if not torch.is_tensor(x) or x.shape[-1] != in_features:
return
ss = x.detach().reshape(-1, in_features).float().square().sum(dim=0)
if rec['ss'] is None:
rec['ss'] = ss
else:
if rec['ss'].device != ss.device: # offload moves blocks between devices mid-run
rec['ss'] = rec['ss'].to(ss.device)
rec['ss'] += ss
rec['n'] += x.numel() // in_features
if rec['n'] >= TOKENS_DONE:
rec['done'] = True
if all(r['done'] for r in capture['recs'].values()):
persist()
return hook
def persist():
"""Write accumulated statistics and stamp them onto the layers.
Runs from the last hook to complete its quota or from the forward
deadline, inside a forward; the write is a few MB once per checkpoint
ever. Layers under the token floor are omitted rather than saved with
meaningless statistics. Handles stay registered but inert until the
next safe point removes them (hook removal here would mutate the hook
dict the forward is iterating).
"""
if capture['complete']:
return
capture['complete'] = True
from safetensors.torch import save_file
tensors, min_n = {}, None
for name, rec in capture['recs'].items():
if rec['ss'] is None or rec['n'] < TOKENS_FLOOR:
continue
rms = (rec['ss'] / rec['n']).sqrt().float().cpu().contiguous().clone()
tensors[name] = rms
rec['m'].sdnq_calib_rms = rms
min_n = rec['n'] if min_n is None else min(min_n, rec['n'])
if not tensors:
log.warning(f'Network calibration: model="{capture["model"]}" no layer reached {TOKENS_FLOOR} tokens; nothing saved')
return
path = calib_file(capture['model'])
try:
os.makedirs(calib_root, exist_ok=True)
save_file(tensors, path, metadata={'version': '1', 'model': capture['model'], 'tokens': str(min_n)})
log.info(f'Network calibration: model="{capture["model"]}" layers={len(tensors)}/{len(capture["recs"])} tokens={min_n} saved="{path}"')
except Exception as e:
log.warning(f'Network calibration: save failed path="{path}" {e}')
def maybe_detach():
"""Remove inert hooks once capture finished; safe only outside a model forward."""
if capture['complete'] and capture['handles']:
detach_capture()
def load_stats(model_name, modules_list):
from safetensors import safe_open
path = calib_file(model_name)
loaded = 0
with safe_open(path, framework='pt', device='cpu') as f:
keys = set(f.keys())
for name, m in modules_list:
if name in keys:
m.sdnq_calib_rms = f.get_tensor(name)
loaded += 1
log.info(f'Network calibration: model="{model_name}" layers={loaded} loaded="{path}"')
def on_model_loaded(sd_model):
detach_capture()
if not enabled():
return
name = checkpoint_name(sd_model)
if name is None:
return
modules_list = eligible_modules(sd_model)
if len(modules_list) == 0:
return
if os.path.isfile(calib_file(name)):
load_stats(name, modules_list)
return
if 'Model' in (getattr(shared.opts, 'cuda_compile', None) or []):
return # hooks inside a compiled module graph-break or misbehave; skip capture entirely
capture['model'] = name
capture['handles'].append(denoiser_root(sd_model).register_forward_pre_hook(deadline_hook))
for mod_name, m in modules_list:
rec = {'m': m, 'ss': None, 'n': 0, 'done': False}
capture['recs'][mod_name] = rec
capture['handles'].append(m.register_forward_pre_hook(hook_for(rec, int(m.sdnq_dequantizer.original_shape[-1]))))
log.info(f'Network calibration: model="{name}" layers={len(modules_list)} collecting activation statistics')
def rms_for(layer):
"""Per-channel input RMS for a layer, or None when absent or disabled."""
maybe_detach()
if not enabled():
return None
return getattr(layer, 'sdnq_calib_rms', None)
script_callbacks.on_model_loaded(on_model_loaded)
modules/lora/lora_common.py
+4 -1
View File
@@ -1,6 +1,6 @@
import os
from modules.lora import lora_timers
from modules.lora import network_lora, network_hada, network_ia3, network_oft, network_lokr, network_full, network_norm, network_glora
from modules.lora import network_lora, network_hada, network_ia3, network_oft, network_boft, network_lokr, network_full, network_norm, network_glora
timer = lora_timers.Timer()
@@ -9,6 +9,7 @@ module_types = [
network_lora.ModuleTypeLora(),
network_hada.ModuleTypeHada(),
network_ia3.ModuleTypeIa3(),
network_boft.ModuleTypeBOFT(), # ahead of oft, which claims any oft_blocks key without checking its rank
network_oft.ModuleTypeOFT(),
network_lokr.ModuleTypeLokr(),
network_full.ModuleTypeFull(),
@@ -18,3 +19,5 @@ module_types = [
loaded_networks: list = [] # no type due to circular import
previously_loaded_networks: list = [] # no type due to circular import
extra_network_lora = None # initialized in extra_networks.py
last_backup_size: int = 0 # bytes of weight backups the last activate pass held
last_mode: str = '' # how that pass left the weights: backup, fuse or factor
modules/lora/lora_convert.py
+14 -6
View File
@@ -136,8 +136,9 @@ class KeyConvert:
sd_module = shared.sd_model.network_layer_mapping.get(flat_key, None)
if sd_module is not None:
key = flat_key
if debug and sd_module is None:
raise RuntimeError(f"LoRA key not found in network_layer_mapping: key={key} mapping={shared.sd_model.network_layer_mapping.keys()}")
if sd_module is None:
if debug:
raise RuntimeError(f"LoRA key not found in network_layer_mapping: key={key} mapping={shared.sd_model.network_layer_mapping.keys()}")
return key, sd_module
@@ -486,28 +487,35 @@ def assign_network_names_to_compvis_modules(sd_model):
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) # wrapped model compatibility
network_layer_mapping = {}
if hasattr(sd_model, 'text_encoder') and sd_model.text_encoder is not None:
for name, module in sd_model.text_encoder.named_modules():
for name, module in sd_model.text_encoder.named_modules() :
prefix = "lora_te1_" if hasattr(sd_model, 'text_encoder_2') else "lora_te_"
network_name = prefix + name.replace(".", "_")
network_layer_mapping[network_name] = module
module.network_layer_name = network_name
if hasattr(sd_model, 'text_encoder_2'):
if hasattr(sd_model, 'text_encoder_2') and sd_model.text_encoder_2 is not None:
for name, module in sd_model.text_encoder_2.named_modules():
network_name = "lora_te2_" + name.replace(".", "_")
network_layer_mapping[network_name] = module
module.network_layer_name = network_name
if hasattr(sd_model, 'unet'):
if hasattr(sd_model, 'unet') and sd_model.unet is not None:
for name, module in sd_model.unet.named_modules():
network_name = "lora_unet_" + name.replace(".", "_")
network_layer_mapping[network_name] = module
module.network_layer_name = network_name
if hasattr(sd_model, 'transformer'):
if hasattr(sd_model, 'transformer') and sd_model.transformer is not None:
for name, module in sd_model.transformer.named_modules():
network_name = "lora_transformer_" + name.replace(".", "_")
network_layer_mapping[network_name] = module
if "norm" in network_name and "linear" not in network_name and shared.sd_model_type != "sd3":
continue
module.network_layer_name = network_name
if hasattr(sd_model, 'transformer_ref') and sd_model.transformer_ref is not None:
for name, module in sd_model.transformer_ref.named_modules():
network_name = "lora_transformer_" + name.replace(".", "_")
network_layer_mapping[network_name] = module
if "norm" in network_name and "linear" not in network_name and shared.sd_model_type != "sd3":
continue
module.network_layer_name = network_name
if hasattr(sd_model, 'llm_adapter') and sd_model.llm_adapter is not None:
for name, module in sd_model.llm_adapter.named_modules():
network_name = "lora_llm_adapter_" + name.replace(".", "_")
modules/lora/lora_diffusers.py
+10 -6
View File
@@ -4,6 +4,7 @@ import diffusers
from modules import shared, errors
from modules.logger import log
from modules.lora import network
from modules.lora import lora_overrides
from modules.lora import lora_common as l
@@ -53,27 +54,30 @@ def load_per_module(sd_model: diffusers.DiffusionPipeline, filename: str, adapte
def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale:float=shared.opts.extra_networks_default_multiplier, lora_module=None, reason: str = '') -> network.Network | None:
t0 = time.time()
name = name.replace(".", "_")
reason = 'unknown' if reason is None or len(reason) == 0 else reason
sd_model: diffusers.DiffusionPipeline = getattr(shared.sd_model, "pipe", shared.sd_model)
log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason={reason or "unknown"} scale={lora_scale} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" detected={network_on_disk.sd_version} method=diffusers reason="{reason}" scale={lora_scale} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
if not hasattr(sd_model, 'load_lora_weights'):
log.error(f'Network load: type=LoRA class={sd_model.__class__} does not implement load lora')
log.error(f'Network load: type=LoRA class={sd_model.__class__} method=diffusersdoes not implement load lora')
return None
try:
if lora_module is not None and isinstance(lora_module, list) and len(lora_module) > 0:
name = load_per_module(sd_model, network_on_disk.filename, adapter_name=name, lora_modules=lora_module)
sd_model._lora_partial = True # pylint: disable=protected-access
else:
if shared.sd_model_type in ['sd', 'sdxl']: # skip te to avoid errors when lora does not have te to start with
diffusers.loaders.lora_pipeline._load_lora_into_text_encoder = lambda *args, **kwargs: None # pylint: disable=protected-access
sd_model.load_lora_weights(network_on_disk.filename, adapter_name=name)
except Exception as e:
if 'already in use' in str(e):
pass
else:
if 'following keys have not been correctly renamed' in str(e):
log.error(f'Network load: type=LoRA name="{name}" diffusers unsupported format')
log.error(f'Network load: type=LoRA name="{name}" method=diffusers unsupported format')
elif 'object has no attribute' in str(e):
log.error(f'Network load: type=LoRA name="{name}" diffusers empty module')
log.error(f'Network load: type=LoRA name="{name}" method=diffusers empty module')
else:
log.error(f'Network load: type=LoRA name="{name}" {e}')
log.error(f'Network load: type=LoRA name="{name}" method=diffusers {e}')
if l.debug:
errors.display(e, "LoRA")
return None
@@ -83,7 +87,7 @@ def load_diffusers(name: str, network_on_disk: network.NetworkOnDisk, lora_scale
list_adapters = sd_model.get_list_adapters()
list_adapters = [adapter for adapters in list_adapters.values() for adapter in adapters]
if name not in list_adapters:
log.error(f'Network load: type=LoRA name="{name}" adapters={list_adapters} not loaded')
log.error(f'Network load: type=LoRA name="{name}" method=diffusers adapters={list_adapters} not loaded')
else:
diffuser_loaded.append(name)
diffuser_scales.append(lora_scale)
modules/lora/lora_factor_cache.py
+250
View File
@@ -0,0 +1,250 @@
"""Disk cache for hosted svd factors.
Hosting a non-factorable adapter set costs one truncated svd per targeted
layer (tens of ms each, seconds per file) every time the set is applied
fresh. The resulting factors are deterministic in the checkpoint, the loaded
set (files, multipliers, dyn_dim), the host rank and the calibration
statistics, so they are cached on disk keyed by exactly that identity and
replayed bit-identically on the next apply of the same configuration.
One safetensors file per configuration under ``models/lora-factor-cache``,
holding every hosted layer's post-rotation factor pair as rowwise int8
with fp32 scales (measured fidelity-free in output space, half the bytes
of bf16). Files are named by the model and network set with an
identity-hash suffix, and the exact signature is embedded in the file
metadata. Factors are quantized before first use: ``store`` returns the
dequantized round-trip for the caller to apply, so a fresh compute and a
later cache hit attach bit-identical tensors. The ``lora_sdnq_host_cache``
option is the size budget in GB (0 disables); least-recently-used entries
are evicted past the budget. Any doubt about identity (unknown checkpoint,
unreadable lora file, signature mismatch) disables caching for the pass
rather than risking a stale hit.
"""
import os
import json
import hashlib
import torch
from modules import paths, shared
from modules.lora import lora_common as l
from modules.logger import log
cache_root = os.path.join(paths.models_path, 'lora-factor-cache')
state = {'wn': None, 'sig': None, 'path': None, 'store': {}, 'dirty': False, 'hits': 0, 'misses': 0}
FMT = '5' # bump on entry-layout changes so older files recompute instead of replaying short
def budget_gb():
try:
return float(getattr(shared.opts, 'lora_sdnq_host_cache', 0) or 0)
except Exception:
return 0.0
def signature(wanted_names):
"""Content identity of a hosted-apply configuration, or None when caching is unsafe."""
from modules.lora import lora_calib
model_name = lora_calib.checkpoint_name(getattr(shared, 'sd_model', None))
if model_name is None:
return None
calib_path = lora_calib.calib_file(model_name)
from modules.lora import lora_stack
parts = {
'model': model_name,
'rank': int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0),
'calib': int(os.path.getmtime(calib_path)) if lora_calib.enabled() and os.path.isfile(calib_path) else None, # the toggle is part of the identity: factors computed under the other setting must not replay
'stack': lora_stack.signature(),
'nets': [],
}
from modules.lora import lora_blocks
for name, te, unet, dyn in wanted_names:
net = next((n for n in l.loaded_networks if n.name == name), None)
filename = getattr(getattr(net, 'network_on_disk', None), 'filename', None)
try:
st = os.stat(filename)
except Exception:
return None
entry = [name, repr(te), repr(unet), repr(dyn), filename, int(st.st_mtime), st.st_size]
spec = lora_blocks.net_signature(net)
if spec is not None: # appended only when set so existing cache files stay valid without block weights
entry.append(spec)
parts['nets'].append(entry)
return parts
def label(parts):
"""Filename prefix from the model and net names, so the cache folder reads without tooling."""
names = [parts['model'].replace('\\', '/').split('/')[-1]] + [n[0] for n in parts['nets']]
text = '-'.join(names)
text = ''.join(c if c.isalnum() or c in '._-' else '-' for c in text)
return text[:96]
def begin_pass(wanted_names):
"""Bind the pass to its cache entry; identity-memoized on the wanted_names tuple."""
if wanted_names is state['wn']:
return
state['wn'] = wanted_names
state.update(sig=None, path=None, dirty=False)
state['store'] = {}
if budget_gb() <= 0 or wanted_names == ():
return
parts = signature(wanted_names)
if parts is None:
return
sig = json.dumps(parts, sort_keys=True)
key = hashlib.sha256(sig.encode()).hexdigest()[:24]
path = os.path.join(cache_root, f'{label(parts)}-{key}.safetensors')
entries = {}
if os.path.isfile(path):
try:
from safetensors import safe_open
with safe_open(path, framework='pt', device='cpu') as f:
meta = f.metadata() or {}
if meta.get('sig') == sig and meta.get('fmt') == FMT:
for k in f.keys():
entries[k] = f.get_tensor(k)
os.utime(path, None) # freshness for LRU eviction
except Exception as e:
log.debug(f'Network cache: read failed path="{path}" {e}')
entries = {}
state.update(sig=sig, path=path)
state['store'] = entries
log.debug(f'Network cache: entry="{path}" keys={len(entries)}')
def quantize_rowwise(t):
t32 = t.detach().to(torch.float32)
scale = t32.abs().amax(dim=1, keepdim=True).clamp(min=1e-12) / 127.0
q = (t32 / scale).round().clamp(-127, 127).to(torch.int8)
return q, scale
def dequantize_rowwise(q, scale):
# int8 * fp32 with a single fp32 rounding: identical on any device, so hit and miss replay the same values
return q.to(torch.float32) * scale
def lookup(network_layer_name):
"""Cached (up, down, energy, calibrated, rms) for a layer, or None; factors return as fp32.
Pure lookup with no hit/miss accounting: the fast-path probe uses it so a
layer is only counted once, by whichever caller consumes the answer.
"""
if state['sig'] is None:
return None
st = state['store']
up_q, up_s = st.get(f'{network_layer_name}.up_q'), st.get(f'{network_layer_name}.up_s')
down_q, down_s = st.get(f'{network_layer_name}.down_q'), st.get(f'{network_layer_name}.down_s')
energy = st.get(f'{network_layer_name}.energy')
calib = st.get(f'{network_layer_name}.calib')
rms = st.get(f'{network_layer_name}.rms')
if up_q is None or up_s is None or down_q is None or down_s is None or energy is None or calib is None or rms is None:
return None
return dequantize_rowwise(up_q, up_s), dequantize_rowwise(down_q, down_s), float(energy), bool(calib), float(rms)
def note_hit():
state['hits'] += 1
def fetch(network_layer_name):
"""``lookup`` with accounting: a usable entry counts a hit, anything else a miss."""
entry = lookup(network_layer_name)
if entry is None:
if state['sig'] is not None:
state['misses'] += 1
return None
state['hits'] += 1
return entry
def lookup_scores(network_layer_name):
"""Cached select scores for a layer as ((s0, s1), (a0, a1)), or None.
Score records ride the same signature-keyed entry as factors, and the
signature already pins everything the scores depend on (pair, multipliers,
stack mode and params). No hit/miss accounting: a record saves scoring and
delta assembly, not a sketch.
"""
if state['sig'] is None:
return None
t = state['store'].get(f'{network_layer_name}.sel')
if t is None:
return None
return (float(t[0]), float(t[1])), (float(t[2]), float(t[3]))
def store_scores(network_layer_name, scores, abs_sums):
"""Persist a select-mode score record; additive to the entry, older files upgrade on their next pass."""
if state['sig'] is None:
return
state['store'][f'{network_layer_name}.sel'] = torch.tensor([scores[0], scores[1], abs_sums[0], abs_sums[1]], dtype=torch.float64)
state['dirty'] = True
def store(network_layer_name, up, down, energy, calibrated, rms):
"""Quantize-before-use: returns the dequantized round-trip the caller must apply.
The factors quantize to rowwise int8 whether or not a cache entry can be
written, so the factors applied now, the factors a later hit replays, and a
cache-off apply are the same tensors (the round-trip also zeroes null-tail
columns the attach-side trim relies on). ``rms`` is the assembled delta's
rms, kept so replays can evaluate the requantize routing rule without
assembling the delta.
"""
up_q, up_s = quantize_rowwise(up)
down_q, down_s = quantize_rowwise(down)
if state['sig'] is not None:
st = state['store']
st[f'{network_layer_name}.up_q'] = up_q.to('cpu').contiguous()
st[f'{network_layer_name}.up_s'] = up_s.to('cpu').contiguous()
st[f'{network_layer_name}.down_q'] = down_q.to('cpu').contiguous()
st[f'{network_layer_name}.down_s'] = down_s.to('cpu').contiguous()
st[f'{network_layer_name}.energy'] = torch.tensor(float(energy))
st[f'{network_layer_name}.calib'] = torch.tensor(1 if calibrated else 0, dtype=torch.uint8)
st[f'{network_layer_name}.rms'] = torch.tensor(float(rms))
state['dirty'] = True
return dequantize_rowwise(up_q, up_s).to(up.dtype), dequantize_rowwise(down_q, down_s).to(down.dtype)
def evict():
budget = budget_gb() * 2**30
try:
files = [os.path.join(cache_root, f) for f in os.listdir(cache_root) if f.endswith('.safetensors')]
sizes = {p: os.path.getsize(p) for p in files}
except Exception:
return
total = sum(sizes.values())
for p in sorted(files, key=os.path.getmtime):
if total <= budget:
break
if p == state['path']:
continue # never evict the entry of the live pass
try:
os.remove(p)
total -= sizes[p]
except Exception:
pass
def flush():
"""Persist a dirty pass store; returns (hits, misses) since the last flush."""
hits, misses = state['hits'], state['misses']
state['hits'] = state['misses'] = 0
if not state['dirty'] or state['path'] is None:
return hits, misses
state['dirty'] = False
try:
from safetensors.torch import save_file
os.makedirs(cache_root, exist_ok=True)
tmp = state['path'] + '.tmp'
save_file(state['store'], tmp, metadata={'sig': state['sig'], 'fmt': FMT})
os.replace(tmp, state['path'])
evict()
except Exception as e:
log.warning(f'Network cache: write failed path="{state["path"]}" {e}')
return hits, misses
modules/lora/lora_load.py
+64 -40
View File
@@ -28,6 +28,7 @@ NATIVE_DISPATCH = {
'f2': 'pipelines.flux.flux2_lora',
'anima': 'pipelines.anima.anima_lora',
'krea2': 'pipelines.krea2.krea2_lora',
'minimaxh3': 'pipelines.minimax.minimax_lora',
}
@@ -55,12 +56,19 @@ def lora_dump(lora, dct):
def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Network | None:
if not shared.sd_loaded:
return None
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
# cached
cached = lora_cache.get(name, None)
if cached is not None:
if l.debug:
log.trace(f'LoRA: load name="{name}" fn="{network_on_disk.filename}" cache=True')
return cached
# native dispatch
native_module = NATIVE_DISPATCH.get(shared.sd_model_type)
if l.debug:
log.trace(f'LoRA: load name="{name}" fn="{network_on_disk.filename}" native={native_module}')
if native_module is not None:
import importlib
mod = importlib.import_module(native_module)
@@ -68,6 +76,10 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
if net is not None:
lora_cache[name] = net
return net
# fallback to standard network loading
if l.debug:
log.trace(f'LoRA: load name="{name}" network_on_disk="{network_on_disk.filename}" safetensors')
net = network.Network(name, network_on_disk)
net.mtime = os.path.getmtime(network_on_disk.filename)
state_dict = sd_models.read_state_dict(network_on_disk.filename, what='network')
@@ -95,7 +107,7 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
emb_dict[vec_name] = weight
bundle_embeddings[emb_name] = emb_dict
continue
if parts[0] in ["clip_l","clip_g","t5","unet","transformer"]:
if parts[0] in ["clip_l", "clip_g", "t5", "unet", "transformer", "transformer_2"]:
network_part = []
while parts and parts[-1] in ["alpha","weight","lora_up","lora_down"]:
network_part.insert(0,parts[-1])
@@ -107,7 +119,6 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
if key_network_without_network_parts.startswith("unet") or key_network_without_network_parts.startswith("transformer"):
key_network_without_network_parts = "lora_" + key_network_without_network_parts
key_network_without_network_parts = key_network_without_network_parts.replace("clip_g","lora_te2").replace("clip_l","lora_te")
# TODO lora: add t5 key support for sd35/f1
elif len(parts) > 5: # messy handler for diffusers peft lora
key_network_without_network_parts = '_'.join(parts[:-2])
@@ -147,9 +158,9 @@ def load_safetensors(name, network_on_disk: network.NetworkOnDisk) -> network.Ne
if len(keys_failed_to_match) > 0:
log.warning(f'Network load: type=LoRA name="{name}" type={set(network_types)} unmatched={len(keys_failed_to_match)} matched={len(matched_networks)}')
if l.debug:
log.debug(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
log.trace(f'Network load: type=LoRA name="{name}" unmatched={keys_failed_to_match}')
else:
log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
log.debug(f'Network load: type=LoRA name="{name}" type={set(network_types)} keys={len(matched_networks)} dtypes={dtypes} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
if len(matched_networks) == 0:
return None
lora_cache[name] = net
@@ -172,26 +183,61 @@ def maybe_recompile_model(names, te_multipliers):
if not recompile_model:
skip_lora_load = True
if len(l.loaded_networks) > 0 and l.debug:
log.debug('Model Compile: Skipping LoRa loading')
log.trace('LoRA: recompile required, skip loading')
return recompile_model, skip_lora_load
else:
recompile_model = True
shared.compiled_model_state.lora_model = []
if l.debug:
log.trace(f'LoRA recompile check: task={sd_models.get_diffusers_task(shared.sd_model)} recompile={recompile_model} load={skip_lora_load}')
if recompile_model:
current_task = sd_models.get_diffusers_task(shared.sd_model)
log.debug(f'Compile: task={current_task} force model reload')
backup_cuda_compile = shared.opts.cuda_compile
backup_scheduler = getattr(sd_model, "scheduler", None)
backup_loaded_loras = getattr(sd_model, "loaded_loras", None) # reload below replaces shared.sd_model with a new pipe object
sd_models.unload_model_weights(op='model')
shared.opts.cuda_compile = []
shared.opts.cuda_compile = ['LoRA'] # if its empty, it will be overridden by set_openvino_overrides() to ['Model'] which is not what we want
sd_models.reload_model_weights(op='model')
shared.sd_model = sd_models.set_diffuser_pipe(shared.sd_model, current_task)
shared.opts.cuda_compile = backup_cuda_compile
new_sd_model = getattr(shared.sd_model, "pipe", shared.sd_model) # scheduler/cache must be reapplied to the new object, not the discarded one
if backup_scheduler is not None:
sd_model.scheduler = backup_scheduler
new_sd_model.scheduler = backup_scheduler
if backup_loaded_loras is not None:
new_sd_model.loaded_loras = backup_loaded_loras
from modules import processing_diffusers # pylint: disable=import-outside-toplevel
processing_diffusers.orig_pipeline = shared.sd_model # otherwise process_diffusers() restores the pre-recompile pipeline once generation ends
return recompile_model, skip_lora_load
def add_network(filename):
"""Register one network file in the available-network tables."""
if not os.path.isfile(filename):
return
name = os.path.splitext(os.path.basename(filename))[0]
name = name.replace('.', '_')
try:
entry = network.NetworkOnDisk(name, filename)
available_networks[entry.name] = entry
if entry.alias in available_network_aliases:
forbidden_network_aliases[entry.alias.lower()] = 1
available_network_aliases[entry.name] = entry
if entry.fullname != entry.name:
available_network_aliases[entry.fullname] = entry
# entry.name mangles dots to underscores for legacy reasons and entry.fullname
# carries any subfolder prefix, so neither matches when the user types the file's
# natural basename. setdefault avoids clobbering an explicit primary entry when
# two files in different subfolders share a basename.
basename_alias = os.path.splitext(os.path.basename(filename))[0]
if basename_alias and basename_alias not in (entry.name, entry.fullname):
available_network_aliases.setdefault(basename_alias, entry)
if entry.shorthash:
available_network_hash_lookup[entry.shorthash] = entry
except OSError as e: # should catch FileNotFoundError and PermissionError etc.
log.error(f'LoRA: filename="{filename}" {e}')
def list_available_networks():
t0 = time.time()
available_networks.clear()
@@ -202,31 +248,6 @@ def list_available_networks():
if not os.path.exists(shared.cmd_opts.lora_dir):
log.warning(f'LoRA directory not found: path="{shared.cmd_opts.lora_dir}"')
def add_network(filename):
if not os.path.isfile(filename):
return
name = os.path.splitext(os.path.basename(filename))[0]
name = name.replace('.', '_')
try:
entry = network.NetworkOnDisk(name, filename)
available_networks[entry.name] = entry
if entry.alias in available_network_aliases:
forbidden_network_aliases[entry.alias.lower()] = 1
available_network_aliases[entry.name] = entry
if entry.fullname != entry.name:
available_network_aliases[entry.fullname] = entry
# entry.name mangles dots to underscores for legacy reasons and entry.fullname
# carries any subfolder prefix, so neither matches when the user types the file's
# natural basename. setdefault avoids clobbering an explicit primary entry when
# two files in different subfolders share a basename.
basename_alias = os.path.splitext(os.path.basename(filename))[0]
if basename_alias and basename_alias not in (entry.name, entry.fullname):
available_network_aliases.setdefault(basename_alias, entry)
if entry.shorthash:
available_network_hash_lookup[entry.shorthash] = entry
except OSError as e: # should catch FileNotFoundError and PermissionError etc.
log.error(f'LoRA: filename="{filename}" {e}')
candidates = sorted(files_cache.list_files(shared.cmd_opts.lora_dir, ext_filter=[".pt", ".ckpt", ".safetensors"]))
with concurrent.futures.ThreadPoolExecutor(max_workers=shared.max_workers) as executor:
for fn in candidates:
@@ -261,7 +282,7 @@ def gather_networks(names):
return networks_on_disk
def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None, activate=True):
def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=None, lora_modules=None, block_specs=None, activate=True):
networks_on_disk = gather_networks(names)
failed_to_load_networks = []
recompile_model, skip_lora_load = maybe_recompile_model(names, te_multipliers)
@@ -277,7 +298,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
if network_on_disk is not None:
shorthash = getattr(network_on_disk, 'shorthash', '').lower()
if l.debug:
log.debug(f'Network load: type=LoRA name="{name}" file="{network_on_disk.filename}" hash="{shorthash}" cached={name in lora_cache}')
log.trace(f'LoRA: name="{name}" fn="{network_on_disk.filename}" hash="{shorthash}" cached={name in lora_cache}')
try:
lora_scale = te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier
lora_module = lora_modules[i] if lora_modules and len(lora_modules) > i else None
@@ -301,6 +322,8 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
if net is None:
failed_to_load_networks.append(name)
lora_ver = network_on_disk.sd_version if network_on_disk is not None else None
if lora_ver is None or len(lora_ver) == 0:
lora_ver = "unknown"
log.error(f'Network load: type=LoRA name="{name}" detected={lora_ver} not loaded')
continue
if hasattr(sd_model, 'embedding_db'):
@@ -309,6 +332,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
'te': te_multipliers[i] if te_multipliers else shared.opts.extra_networks_default_multiplier,
'unet': unet_multipliers[i] if unet_multipliers else shared.opts.extra_networks_default_multiplier,
'dyn': dyn_dims[i] if dyn_dims else None, # a multiplier is not a rank; float dyn_dim crashes every consumer that slices with it
'blocks': block_specs[i] if block_specs and len(block_specs) > i else None,
}
l.loaded_networks.append(net)
@@ -321,8 +345,8 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
try:
t1 = time.time()
if l.debug:
log.trace(f'Network load: type=LoRA list={sd_model.get_list_adapters()}')
log.trace(f'Network load: type=LoRA active={sd_model.get_active_adapters()}')
log.trace(f'LoRA: list={sd_model.get_list_adapters()}')
log.trace(f'LoRA: active={sd_model.get_active_adapters()}')
sd_model.set_adapters(adapter_names=lora_diffusers.diffuser_loaded, adapter_weights=lora_diffusers.diffuser_scales)
sd_model.enable_lora() # set_adapters does not clear the disabled flag left by a prior removal
except Exception as e:
@@ -333,7 +357,7 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
try:
if shared.opts.lora_fuse_diffusers and not lora_overrides.disable_fuse():
sd_model.fuse_lora(adapter_names=lora_diffusers.diffuser_loaded, lora_scale=1.0, fuse_unet=True, fuse_text_encoder=True) # diffusers with fuse uses fixed scale since later apply does the scaling
sd_model.unload_lora_weights()
# sd_model.unload_lora_weights() # optionally unload fused lora as we dont need it, but it may cause issues with some models
l.timer.activate += time.time() - t1
except Exception as e:
log.error(f'Network load: type=LoRA action=fuse {str(e)}')
@@ -350,10 +374,10 @@ def network_load(names, te_multipliers=None, unet_multipliers=None, dyn_dims=Non
networks.network_activate()
if len(l.loaded_networks) > 0 and l.debug:
log.debug(f'Network load: type=LoRA loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={shared.opts.lora_fuse_native}:{shared.opts.lora_fuse_diffusers}')
log.trace(f'LoRA: loaded={[n.name for n in l.loaded_networks]} cache={list(lora_cache)} fuse={lora_overrides.fuse_native()}:{shared.opts.lora_fuse_diffusers}')
if recompile_model:
log.info("Network load: type=LoRA recompiling model")
log.info("Network load: type=LoRA model recompile required")
if shared.compiled_model_state is not None:
backup_lora_model = shared.compiled_model_state.lora_model
else:
modules/lora/lora_nunchaku.py
+9 -2
View File
@@ -1,12 +1,19 @@
import time
from modules import shared, errors
from modules.logger import log
from modules.lora import lora_load, lora_common
from modules.lora import lora_load, lora_common, network
previously_loaded = [] # we maintain private state here
def wrap_network(network_on_disk):
net = network.Network(network_on_disk.name, network_on_disk)
net.mentioned_name = network_on_disk.name
network_on_disk.read_hash() # nothing else on this path fills the hash infotext reads
return net
def load_nunchaku(names, strengths):
global previously_loaded # pylint: disable=global-statement
strengths = [s[0] if isinstance(s, list) else s for s in strengths]
@@ -26,7 +33,7 @@ def load_nunchaku(names, strengths):
from nunchaku.lora.flux.compose import compose_lora
composed_lora = compose_lora(loras)
shared.sd_model.transformer.update_lora_params(composed_lora)
lora_common.loaded_networks = [n[0] for n in networks] # used by infotext
lora_common.loaded_networks[:] = [wrap_network(n[0]) for n in networks] # read by infotext and the trigger tags
t1 = time.time()
lora_common.timer.load = t1 - t0
log.debug(f"Network load: type=LoRA method=nunchaku loras={names} strength={strengths} time={t1-t0:.3f}")
modules/lora/lora_overrides.py
+68 -5
View File
@@ -1,4 +1,9 @@
import os
from modules import shared
from modules.logger import log
debug_log = log.trace if os.environ.get('SD_LORA_DEBUG', None) is not None else lambda *args, **kwargs: None
force_hashes_diffusers = [ # forced always
@@ -32,6 +37,7 @@ allow_native = [
'anima',
'ernieimage',
'krea2',
'minimaxh3',
]
@@ -77,13 +83,70 @@ def get_method(shorthash=''):
return 'native', 'default'
# Roles a LoRA is fused into; a quantized component in any of them makes fusing unsafe.
fuse_roots = ('transformer', 'unet', 'text_encoder', 'llm_adapter')
def fuse_components(sd_model):
"""Component names a network fuses into, matched by role prefix so numbered and reference siblings are covered."""
names = getattr(sd_model, 'components', None)
if not isinstance(names, dict):
names = vars(sd_model)
return [name for name in names if name.startswith(fuse_roots)]
def is_quantized(module):
"""Return True when ``module`` carries a quantization config.
``config.quantization_config`` is read first: SDNQ sets both it and the plain
attribute when it quantizes in place, but a checkpoint that ships pre-quantized
only reaches the plain attribute through the diffusers ConfigMixin name proxy,
which is deprecated for removal.
"""
if module is None:
return False
config = getattr(module, 'config', None)
if config is not None and getattr(config, 'quantization_config', None) is not None:
return True
return getattr(module, 'quantization_config', None) is not None
def disable_fuse():
if hasattr(shared.sd_model, 'quantization_config'):
"""Return True when fusing a network into model weights is unsafe.
Fusing keeps no pristine copy of the weight, so each apply and restore
round-trips it through its storage format. On quantized weights that is a
dequantize-add-requantize cycle per network swap whose error compounds.
"""
from modules.lora import lora_common as l
from modules.lora import lora_stack
if lora_stack.select_possible(len(l.loaded_networks)) or lora_stack.select_engaged():
debug_log('LoRA: fuse=False reason="active select mode"')
return True # select flips per-layer winners against the pristine backup; a dormant select mode leaves fuse alone
sd_model = getattr(shared.sd_model, 'pipe', shared.sd_model)
if is_quantized(sd_model):
debug_log('LoRA: fuse=False reason="model is quantized"')
return True
if hasattr(shared.sd_model, 'transformer') and hasattr(shared.sd_model.transformer, 'quantization_config'):
if any(is_quantized(getattr(sd_model, name, None)) for name in fuse_components(sd_model)):
debug_log('LoRA: fuse=False reason="component is quantized"')
return True
if hasattr(shared.sd_model, 'transformer_2') and hasattr(shared.sd_model.transformer_2, 'quantization_config'):
if hasattr(sd_model, '_lora_partial'):
debug_log('LoRA: fuse=False reason="partial lora applied"')
return True
if hasattr(shared.sd_model, '_lora_partial'):
if shared.sd_model_type in fuse_ignore:
debug_log(f'LoRA: fuse=False reason="model type {shared.sd_model_type} in fuse_ignore"')
return True
return shared.sd_model_type in fuse_ignore
return False
def fuse_native():
"""Return True when the native apply path may fuse into model weights.
The single source of truth for the native fuse decision: it must agree across
the backup, activate and deactivate passes, since backup mode restores from a
stored tensor while fuse mode restores by subtracting the delta.
"""
result = shared.opts.lora_fuse_native and not disable_fuse()
force = os.environ.get('SD_LORA_FUSE', None) is not None
debug_log(f'LoRA: native fuse={result} force={force}')
return (result or force)
modules/lora/lora_sdnq.py
+625
View File
@@ -0,0 +1,625 @@
"""Exact LoRA application for SDNQ-quantized layers.
Baking a LoRA into a quantized weight requantizes it: dequantize, add the
delta, re-round onto the integer grid. When the per-element delta is smaller
than half a quantization step (a rank-decomposed delta on a uint4 layer sits
at a few percent of a step), rounding erases it; what survives is the two
grid-extrema elements per quantization group (2/group_size of the signal)
plus grid-shift noise of the same norm as the delta. The optimal in-grid
representation provably retains ~0%, so no rewrite of the stored integers
can fix this.
The exact path instead rides the SDNQ svd side-channel: the dequantizer
computes ``W = dq(q) + svd_up @ svd_down`` in the rotated domain at full
precision, in every forward mode. A LoRA delta ``B @ A`` is appended as
extra columns of ``svd_up`` and rows of ``svd_down``; because the Hadamard
rotation is block-diagonal, symmetric and self-inverse, storing ``A·H`` for
the down factor makes the round trip exact: ``(B @ (A·H)) · H = B @ A``.
Quantized weights are never touched, so apply and remove are exact and no
weight backup is needed. The side-channel storage is lossless; realized
fidelity floors at the compute dtype, because the dequantizer materializes
``base + factors`` in the result dtype and a delta below its ULP of the
base rounds exactly as it would on an unquantized model of that dtype.
Only additive low-rank modules ride the channel exactly (plain LoRA: no
DoRA, no CP ``mid``, no LyCORIS dense-bias, no ``diff_b``). On sub-8-bit
formats, sets with non-factorable contributions are hosted instead: the
families' own ``calc_updown`` delta is truncated to its top singular
directions and appended the same way, stored at the delta's effective
rank when the spectrum ends in a numerically null tail (dense-combined
plain pairs, low-rank LyCORIS). Truncation keeps the dominant part
of the effect and drops an orthogonal residual, where requantize keeps
only the grid extrema and adds grid-shift noise of the delta's own
magnitude. When activation statistics for the checkpoint exist (see
``lora_calib``), the truncation is channel-weighted to minimize output
error instead of weight error. At 8 bits and above requantize retains
most of the delta, so hosting is skipped there and the requantize path
remains.
A small tail of deltas inverts the tradeoff: when the delta is large
against the grid step AND the truncation genuinely cuts it, requantize
retains more than hosting drops, and the layer routes back to the
requantize path (``REQUANT_RATIO``/``REQUANT_ENERGY``). Both terms must
agree: a thin delta rounds away on the grid however low its capture, and
a low-rank delta hosts exactly however fat it is.
"""
import torch
from modules import devices, shared
from modules.lora import lora_calib, lora_factor_cache, lora_stack # lora_calib registers its model-load hook on import, so this one has to stay eager
from modules.lora import lora_common as l
from modules.logger import log
fallback_layers: list[str] = []
hosted_layers: list[tuple[str, float, bool]] = []
hosted_ranks: list[int] = []
factor_layers: list[str] = []
select_layers: list[str] = []
routed_layers: list[str] = []
REQUANT_RATIO = 0.30 # delta rms over mean grid step above which requantize can retain the delta
REQUANT_ENERGY = 0.90 # sketch capture below which truncation genuinely loses part of it
NULL_TAIL_EPS = 1e-6 # spectrum tail below this fraction of the capture is numerically null; dropping it keeps stored rank at the delta's effective rank
def rank_bucket(r):
"""Fixed rank ladder for compiled-graph reuse: powers of two up to 256, multiples of 64 above (hosted rank plus exact members)."""
if r <= 8:
return 8
if r <= 256:
return 1 << (r - 1).bit_length()
return -(-r // 64) * 64
def pad_rank(t, dim, bucket):
if t.shape[dim] >= bucket:
return t
shape = list(t.shape)
shape[dim] = bucket - t.shape[dim]
return torch.cat([t, t.new_zeros(shape)], dim=dim)
def enabled():
"""True while the exact svd-channel machinery may take quantized layers; the requantize choice routes every layer to the legacy weight-rewrite path."""
return getattr(shared.opts, 'lora_sdnq_apply', 'exact') != 'requantize'
def signature():
"""Identity suffix for the per-module apply stamp; empty on the default exact mechanism."""
return '' if enabled() else '|quant=requantize'
def trim_null_tail(up_h, down_h):
"""Cache entries stored before tail slicing carry null ranks as exact zero columns; trim to the effective rank on attach."""
nz = (up_h != 0).any(dim=0)
if not bool(nz.all()):
k = max(1, int(nz.nonzero().max().item()) + 1) if bool(nz.any()) else 1
if k < up_h.shape[1]:
return up_h[:, :k].contiguous(), down_h[:k].contiguous()
return up_h, down_h
def get_module_factors(module, device, dtype, original_shape=None):
"""Return ``(up_eff, down)`` reproducing ``calc_updown`` exactly, or None.
``updown = up @ down * calc_scale() * multiplier()`` for a plain linear
LoRA; the scalars fold into the up factor. ``dyn_dim`` slices ranks the
same way ``lyco_helpers.rebuild_conventional`` does.
"""
if module.__class__.__name__ != 'NetworkModuleLora':
return None
if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None:
return None
if getattr(module, 'mid_model', None) is not None:
return None
up = module.up_model.weight
down = module.down_model.weight
if up.ndim != 2 or down.ndim != 2:
return None
if original_shape is not None and (up.shape[0] != original_shape[0] or down.shape[1] != original_shape[-1]):
return None # factor_candidate skips shape checks for layers already in factor mode; recheck here so a malformed stack falls back instead of raising in cat
dyn_dim = module.network.dyn_dim
if dyn_dim is not None and up.shape[1] != dyn_dim:
up = up[:, :dyn_dim]
down = down[:dyn_dim]
scalar = module.calc_scale() * module.multiplier()
up_eff = up.to(device=device, dtype=torch.float32) * scalar
return up_eff.to(dtype=dtype), down.to(device=device, dtype=dtype)
def factor_candidate(self, network_layer_name, wanted_names):
"""True when this layer should take the exact svd-append path.
Requires an SDNQ linear layer whose active networks all contribute plain
factorable LoRA modules for this layer. An empty ``wanted_names`` is a
removal request and qualifies whenever factors are currently attached.
"""
if not enabled():
return False # declined layers with factors still attached are stripped by the activate fallthrough
if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear':
return False
if wanted_names != () and lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te'):
if sum(1 for net in l.loaded_networks if net.modules.get(network_layer_name, None) is not None) >= 2:
return False # dense stack modes combine dense deltas; the factor concat would sum
if hasattr(self, 'sdnq_lora_svd_stash'):
return True
if wanted_names == (): # nothing attached, nothing to remove
return False
seen = False
for net in l.loaded_networks:
module = net.modules.get(network_layer_name, None)
if module is None:
continue
seen = True
if module.__class__.__name__ != 'NetworkModuleLora':
return False
if module.dora_scale is not None or module.bias is not None or module.ex_bias is not None or getattr(module, 'mid_model', None) is not None:
return False
if module.up_model.weight.ndim != 2 or module.down_model.weight.ndim != 2:
return False
if module.up_model.weight.shape[0] != self.sdnq_dequantizer.original_shape[0] or module.down_model.weight.shape[1] != self.sdnq_dequantizer.original_shape[-1]:
return False
return seen
def remove_factors(self):
"""Restore the layer's original svd factors; True when factors were attached."""
stash = getattr(self, 'sdnq_lora_svd_stash', None)
if stash is None:
return False
svd_up, svd_down = stash
device = self.scale.device # the stash tuple does not follow module device moves; restore onto wherever the layer lives now
if svd_up is not None and svd_up.device != device:
svd_up = torch.nn.Parameter(svd_up.to(device=device), requires_grad=False)
svd_down = torch.nn.Parameter(svd_down.to(device=device), requires_grad=False)
self.svd_up = svd_up
self.svd_down = svd_down
del self.sdnq_lora_svd_stash
lora_stack.drop(getattr(self, 'network_layer_name', None)) # a selection schedule must not outlive the segments it points into
return True
def apply_factors(self, network_layer_name, wanted_names):
"""Attach the active networks' LoRA factors to this layer's svd side-channel.
Replaces any previously attached factors (multiplier changes re-enter
here with a new ``wanted_names`` signature). Returns True when the layer
changed. Falls back to the caller's requantize path by returning None
when factor extraction fails at this stage.
"""
from sdnq.quant_utils import rotate_hadamard
changed = remove_factors(self)
if wanted_names == ():
return changed
deq = self.sdnq_dequantizer
dtype = deq.result_dtype
ups, downs = [], []
for net in l.loaded_networks:
module = net.modules.get(network_layer_name, None)
if module is None:
continue
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is None:
return None
up_eff, down = factors
if deq.use_hadamard:
down = rotate_hadamard(down.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
ups.append(up_eff)
downs.append(down)
if not ups:
return changed
append_factors(self, ups, downs)
factor_layers.append(network_layer_name)
return True
def append_factors(self, ups, downs):
"""Concatenate ``[out, r]`` / ``[r, in]`` factor pairs onto the layer's svd channel and stash the originals.
Returns the appended parts' rank ranges plus the transposed-layout flag; the
checkpoint's own factors occupy the range before the first entry and bucket
padding lands after the last, so the ranges stay valid on the live buffers.
"""
deq = self.sdnq_dequantizer
device = self.scale.device
dtype = deq.result_dtype
orig_up, orig_down = self.svd_up, self.svd_down
orig_rank = 0
if orig_up is not None:
orig_rank = orig_up.shape[0] if deq.use_quantized_matmul else orig_up.shape[1]
segments, offset = [], orig_rank
for u in ups:
segments.append((offset, offset + u.shape[1]))
offset += u.shape[1]
if deq.use_quantized_matmul:
# matmul layout stores factors transposed: svd_up [r, out], svd_down [in, r]
parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + [u.t() for u in ups]
parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + [d.t() for d in downs]
new_up = torch.cat(parts_up, dim=0).contiguous()
new_down = torch.cat(parts_down, dim=1).contiguous()
else:
parts_up = ([orig_up.to(device=devices.device, dtype=dtype)] if orig_up is not None else []) + ups
parts_down = ([orig_down.to(device=devices.device, dtype=dtype)] if orig_down is not None else []) + downs
new_up = torch.cat(parts_up, dim=1).contiguous()
new_down = torch.cat(parts_down, dim=0).contiguous()
from sdnq.common import use_torch_compile
if use_torch_compile:
# the compiled dequant specializes per factor rank; pad to a fixed bucket so set switches inside a bucket reuse the graph (zero columns contribute exactly nothing)
dim_up, dim_down = (0, 1) if deq.use_quantized_matmul else (1, 0)
bucket = rank_bucket(new_up.shape[dim_up])
new_up = pad_rank(new_up, dim_up, bucket)
new_down = pad_rank(new_down, dim_down, bucket)
self.sdnq_lora_svd_stash = (orig_up, orig_down)
self.svd_up = torch.nn.Parameter(new_up.to(device=device), requires_grad=False)
self.svd_down = torch.nn.Parameter(new_down.to(device=device), requires_grad=False)
return segments, deq.use_quantized_matmul
def channel_candidate(self, network_layer_name, wanted_names):
"""True when this layer can carry a set on the svd channel: quantized, covered, and given a rank to spend."""
if not enabled():
return False
if int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0) <= 0:
return False
if getattr(self, 'sdnq_dequantizer', None) is None or self.__class__.__name__ != 'SDNQLinear':
return False
if wanted_names == ():
return False
return any(net.modules.get(network_layer_name, None) is not None for net in l.loaded_networks)
def select_candidate(self, network_layer_name, wanted_names):
"""True when a select pair can ride this layer's svd channel; pairs ride it at any bit width."""
return channel_candidate(self, network_layer_name, wanted_names)
def host_candidate(self, network_layer_name, wanted_names):
"""True when this layer's set should ride the svd channel as a truncated svd: non-factorable sets below 8 bits, dense-combined sets at any width."""
if not channel_candidate(self, network_layer_name, wanted_names):
return False
if lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te'):
if sum(1 for net in l.loaded_networks if net.modules.get(network_layer_name, None) is not None) >= 2:
return True # combined deltas host at any width: requantizing them is checkpoint-fragile, while single-adapter requantize is well retained
from sdnq.common import dtype_dict
if dtype_dict[self.sdnq_dequantizer.weights_dtype]['num_bits'] >= 8:
return False # requantize retains most of a single set's delta at 8 bits and above; truncation would lose more than it saves
return True
def grid_step(self):
"""Mean grid step in weight units; a codebook layer keeps its Lloyd levels in the scale slot, so its step is their mean adjacent gap."""
scale = self.scale.detach().float()
if self.sdnq_dequantizer.use_codebook:
return float(scale.diff(dim=-1).mean())
return float(scale.mean())
def apply_cached(self, network_layer_name, wanted_names):
"""Attach a hosted set straight from the factor cache, before the delta exists.
Probed by the walk ahead of delta assembly: on a usable entry the routing
rule is evaluated from the stored delta rms and the cached factors attach
exactly as a fetch inside ``apply_hosted`` would, so the pass skips
``calc_updown`` for the layer entirely. Returns True when the layer was
served; None sends the caller down the assemble-and-host path (no entry,
or the rule wants the grid).
"""
from sdnq.quant_utils import rotate_hadamard
lora_factor_cache.begin_pass(wanted_names)
entry = lora_factor_cache.lookup(network_layer_name)
if entry is None:
return None
up_h, down_h, energy, calibrated, rms = entry
up_h, down_h = trim_null_tail(up_h, down_h)
deq = self.sdnq_dequantizer
dtype = deq.result_dtype
remove_factors(self) # before the rule: the svd-channel check must see the checkpoint's own state, and a declined layer must fall through pristine
stack_dense = lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te')
members = []
if not stack_dense:
for net in l.loaded_networks:
module = net.modules.get(network_layer_name, None)
if module is None:
continue
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is not None:
members.append(factors)
if not stack_dense and len(members) == 0 and self.svd_up is None:
step = grid_step(self)
if step > 0 and rms / step > REQUANT_RATIO and energy < REQUANT_ENERGY:
return None # routed to the grid: the caller assembles the delta and requantizes
ups, downs = [], []
for up_eff, down in members:
if deq.use_hadamard:
down = rotate_hadamard(down.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
ups.append(up_eff)
downs.append(down)
lora_factor_cache.note_hit()
append_factors(self, ups + [up_h.to(device=devices.device, dtype=dtype)], downs + [down_h.to(device=devices.device, dtype=dtype)])
hosted_layers.append((network_layer_name, energy, calibrated))
hosted_ranks.append(int(up_h.shape[1]))
return True
def apply_hosted(self, network_layer_name, updown, wanted_names):
"""Host a set's delta on the svd channel: exact factors for factorable
members, the top-k singular directions of the remainder for the rest.
The delta comes from the families' own ``calc_updown``, so every family
and scaling quirk is included; factorable members are subtracted out and
appended exactly so they never compete with the hosted remainder for
rank. When per-checkpoint activation statistics exist (``lora_calib``),
input channels are weighted by their RMS before truncation so the kept
directions minimize output error rather than weight error. Computed
factors are disk-cached per configuration (``lora_factor_cache``) and
replayed bit-identically on later applies. Returns None when the delta
cannot ride the channel (wrong shape) or when the routing rule prefers
the grid for it; the caller falls back to requantize.
"""
from sdnq.quant_utils import rotate_hadamard
deq = self.sdnq_dequantizer
changed = remove_factors(self)
if wanted_names == ():
return changed
if updown is None or updown.ndim != 2 or tuple(updown.shape) != tuple(deq.original_shape):
return None
dtype = deq.result_dtype
members = []
stack_dense = lora_stack.mode() in lora_stack.DENSE_MODES and not network_layer_name.startswith('lora_te')
if not stack_dense: # dense stack modes host the combined delta wholesale; the members' content is already inside it
for net in l.loaded_networks:
module = net.modules.get(network_layer_name, None)
if module is None:
continue
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is not None:
members.append(factors)
# requantize keeps a delta the grid can resolve and that truncation would genuinely
# cut: both terms must agree, since a thin delta rounds away on the grid however
# low its capture, and a low-rank delta hosts exactly however fat it is. Scoped to
# sets the side-channel would otherwise carry whole: factorable members ride
# exactly and dense-combined deltas stay hosted at any magnitude.
delta_rms = float(updown.detach().float().square().mean().sqrt())
maybe_requant = not stack_dense and len(members) == 0 and self.svd_up is None
if maybe_requant:
step = grid_step(self)
maybe_requant = step > 0 and delta_rms / step > REQUANT_RATIO
lora_factor_cache.begin_pass(wanted_names)
cached = lora_factor_cache.fetch(network_layer_name)
D = None if cached is not None else updown.detach().to(devices.device, torch.float32)
ups, downs = [], []
for up_eff, down in members:
if D is not None:
D = D.sub_(up_eff.to(torch.float32) @ down.to(torch.float32)) # factorable members ride exactly; host only the remainder
if deq.use_hadamard:
down = rotate_hadamard(down.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
ups.append(up_eff)
downs.append(down)
if cached is not None:
up_h, down_h, energy, calibrated, _cached_rms = cached
if maybe_requant and energy < REQUANT_ENERGY:
routed_layers.append(network_layer_name)
return None
up_h, down_h = trim_null_tail(up_h, down_h)
append_factors(self, ups + [up_h.to(device=devices.device, dtype=dtype)], downs + [down_h.to(device=devices.device, dtype=dtype)])
hosted_layers.append((network_layer_name, energy, calibrated))
hosted_ranks.append(int(up_h.shape[1]))
return True
up_h, down_h, energy, calibrated = truncate_delta(self, D, dtype)
up_h, down_h = lora_factor_cache.store(network_layer_name, up_h, down_h, energy, calibrated, delta_rms)
if maybe_requant and energy < REQUANT_ENERGY:
routed_layers.append(network_layer_name) # the stored entry memoizes the routing; replays skip the sketch
return None
up_h, down_h = trim_null_tail(up_h, down_h) # the int8 roundtrip zeroes the numeric tail the eps slice keeps; fresh and replayed attaches must trim alike
append_factors(self, ups + [up_h], downs + [down_h])
hosted_layers.append((network_layer_name, energy, calibrated))
hosted_ranks.append(int(up_h.shape[1]))
return True
def truncate_delta(self, D, dtype):
"""Truncate one dense fp32 delta to hosted factors in the layer's channel layout; consumes ``D``.
Calibration-weighted when statistics exist; the sketch is oversampled past
the kept rank so the truncation sits within noise of exact svd. Returns
``(up_h, down_h, energy, calibrated)`` with the down factor rotated into the
layer's hadamard domain.
"""
from sdnq.quant_utils import rotate_hadamard
deq = self.sdnq_dequantizer
cap = int(shared.opts.lora_sdnq_host_rank)
q = min(cap, *D.shape)
rms = lora_calib.rms_for(self)
if rms is not None and rms.shape[-1] == D.shape[-1]:
# scale input channels by their activation RMS so truncation minimizes output error rather than weight error
rms = rms.to(device=D.device, dtype=torch.float32).clamp(min=1e-8)
D = D.mul_(rms)
else:
rms = None
# svd_lowrank draws random projections; fork so user generation seeds are untouched and re-applies are deterministic
with torch.random.fork_rng(devices=[D.device] if D.device.type == 'cuda' else []):
torch.manual_seed(0)
# oversampled sketch with extra power iterations lands within noise of exact svd; only the top q columns are kept
U, S, V = torch.svd_lowrank(D, q=min(q + 64, *D.shape), niter=8)
U, S, V = U[:, :q], S[:q], V[:, :q]
e = S.square()
total_e = e.sum()
if float(total_e) > 0:
# an exactly low-rank delta (dense-combined plain pairs, low-rank LyCORIS) fills the tail with
# numerical zeros; storing them would pad the channel to the cap for nothing
k = int((torch.cumsum(e, 0) < (1.0 - NULL_TAIL_EPS) * total_e).sum().item()) + 1
if k < q:
U, S, V = U[:, :k], S[:k], V[:, :k]
energy = float(S.square().sum() / D.square().sum().clamp(min=1e-30)) # captured fraction, in the weighted domain when calibrated
up_h = (U * S).to(dtype=dtype)
down_h = V.t()
if rms is not None:
down_h = down_h / rms # unscale in the original input basis, before any rotation
if deq.use_hadamard:
down_h = rotate_hadamard(down_h, group_size=deq.hadamard_group_size)
down_h = down_h.to(dtype=dtype)
return up_h, down_h, energy, rms is not None
def apply_select_cached(self, network_layer_name, wanted_names):
"""Serve a select pair from cache and live factors before the walk assembles deltas.
A cached score record plus a factor pair per network (exact factors for
factorable members, cached truncations otherwise) rebuild the segments and
the selection registration without any ``calc_updown``. Returns None when
any piece is missing; the caller assembles and ``apply_select`` recomputes
and stores.
"""
from sdnq.quant_utils import rotate_hadamard
deq = self.sdnq_dequantizer
changed = remove_factors(self)
if wanted_names == ():
return changed
if len(l.loaded_networks) != 2:
return None
dtype = deq.result_dtype
lora_factor_cache.begin_pass(wanted_names)
rec = lora_factor_cache.lookup_scores(network_layer_name)
if rec is None:
return None
pairs, notes = [], []
for i, net in enumerate(l.loaded_networks):
module = net.modules.get(network_layer_name, None)
if module is None:
return None
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is not None:
up_i, down_i = factors
if deq.use_hadamard:
down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
else:
cached = lora_factor_cache.lookup(f'{network_layer_name}#{i}')
if cached is None:
return None
up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype)
notes.append((f'{network_layer_name}#{i}', cached[2], cached[3]))
pairs.append((up_i, down_i))
scores, abs_sums = rec
segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]])
lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums)
for note in notes:
lora_factor_cache.note_hit()
hosted_layers.append(note)
select_layers.append(network_layer_name)
return True
def apply_select(self, network_layer_name, per_net, wanted_names):
"""Attach two networks' contributions as separate side-channel segments for per-layer selection.
Factorable members ride exactly; the rest host as their own truncated svd
with per-net cache entries. Segment ranges and selection scores register
with ``lora_stack``; the flip schedule executes from the step callback.
Returns None when the pair cannot ride the channel; the caller falls back.
"""
from sdnq.quant_utils import rotate_hadamard
deq = self.sdnq_dequantizer
changed = remove_factors(self)
if wanted_names == ():
return changed
if per_net is None or len(per_net) != 2:
return None
dtype = deq.result_dtype
lora_factor_cache.begin_pass(wanted_names)
pairs, ranks = [], []
for i, (net_name, D) in enumerate(per_net):
if D is None or D.ndim != 2 or tuple(D.shape) != tuple(deq.original_shape):
return None
net = next((n for n in l.loaded_networks if n.name == net_name), None)
module = net.modules.get(network_layer_name, None) if net is not None else None
if module is None:
return None
ranks.append(int(getattr(module, 'dim', 0) or 0) or min(int(shared.opts.lora_sdnq_host_rank), *deq.original_shape))
factors = get_module_factors(module, devices.device, dtype, original_shape=deq.original_shape)
if factors is not None:
up_i, down_i = factors
if deq.use_hadamard:
down_i = rotate_hadamard(down_i.to(dtype=torch.float32), group_size=deq.hadamard_group_size).to(dtype=dtype)
else:
key = f'{network_layer_name}#{i}'
cached = lora_factor_cache.fetch(key)
if cached is not None:
up_i, down_i = cached[0].to(device=devices.device, dtype=dtype), cached[1].to(device=devices.device, dtype=dtype)
hosted_layers.append((key, cached[2], cached[3]))
else:
up_i, down_i, energy, calibrated = truncate_delta(self, D.detach().to(devices.device, torch.float32), dtype)
up_i, down_i = lora_factor_cache.store(key, up_i, down_i, energy, calibrated, float(D.detach().float().square().mean().sqrt()))
hosted_layers.append((key, energy, calibrated))
pairs.append((up_i, down_i))
scores, abs_sums = lora_stack.score_pair(per_net[0][1].detach(), per_net[1][1].detach(), ranks[0], ranks[1])
lora_factor_cache.store_scores(network_layer_name, scores, abs_sums)
segments, transposed = append_factors(self, [pairs[0][0], pairs[1][0]], [pairs[0][1], pairs[1][1]])
lora_stack.register(network_layer_name, self, 'factor', scores, segments=(segments[0], segments[1], transposed), abs_sums=abs_sums)
select_layers.append(network_layer_name) # counted apart from the plain concat: both ride the svd channel but only one is a summed set
return True
def note_fallback(self, network_layer_name):
"""Record a quantized layer taking the requantize path (summary-logged per pass); layers the routing rule sent there are counted apart."""
if getattr(self, 'sdnq_dequantizer', None) is not None and network_layer_name not in routed_layers:
fallback_layers.append(network_layer_name)
def reset_pass():
"""Clear every per-pass accumulator, so a pass that raised leaves nothing behind for the next one."""
fallback_layers.clear()
hosted_layers.clear()
hosted_ranks.clear()
factor_layers.clear()
select_layers.clear()
routed_layers.clear() # note_fallback reads this to suppress double counting, so a stale entry silences a real fallback
def report_fallbacks():
hits, misses = lora_factor_cache.flush()
if hits > 0 or misses > 0:
log.info(f'Network load: type=LoRA quant=sdnq cache hits={hits} misses={misses}')
if len(factor_layers) > 0:
log.info(f'Network load: type=LoRA quant=sdnq apply=exact layers={len(factor_layers)}')
factor_layers.clear()
if len(select_layers) > 0:
log.info(f'Network load: type=LoRA quant=sdnq apply=select layers={len(select_layers)} mode={lora_stack.mode()}')
select_layers.clear()
if len(hosted_layers) > 0:
energies = sorted(e for _name, e, _c in hosted_layers)
median = energies[len(energies) // 2]
calibrated = sum(1 for _name, _e, c in hosted_layers if c)
ranks = ''
if len(hosted_ranks) > 0 and min(hosted_ranks) < int(shared.opts.lora_sdnq_host_rank):
rs = sorted(hosted_ranks)
ranks = f' k={rs[0]}-{rs[len(rs) // 2]}-{rs[-1]}' # realized rank spread; shown only when a spectrum collapsed below the cap
log.info(f'Network load: type=LoRA quant=sdnq apply=hosted layers={len(hosted_layers)} rank={int(shared.opts.lora_sdnq_host_rank)}{ranks}{f" calib={calibrated}" if calibrated else ""} energy={median:.2f} min={energies[0]:.2f}')
if l.debug:
log.debug(f'Network load: type=LoRA quant=sdnq hosted={[(n, round(e, 3)) for n, e, _c in hosted_layers[:8]]}{"..." if len(hosted_layers) > 8 else ""}')
hosted_layers.clear()
hosted_ranks.clear()
if len(routed_layers) > 0:
log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(routed_layers)} routed=fat-delta')
if l.debug:
log.debug(f'Network load: type=LoRA quant=sdnq routed={routed_layers[:8]}{"..." if len(routed_layers) > 8 else ""}')
routed_layers.clear()
if len(fallback_layers) > 0:
if enabled():
log.warning(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(fallback_layers)} fidelity=reduced')
else:
log.info(f'Network load: type=LoRA quant=sdnq apply=requantize layers={len(fallback_layers)} reason=setting')
if l.debug:
log.debug(f'Network load: type=LoRA quant=sdnq requantized={fallback_layers[:8]}{"..." if len(fallback_layers) > 8 else ""}')
fallback_layers.clear()
modules/lora/lora_stack.py
+418
View File
@@ -0,0 +1,418 @@
"""Stack modes for combining multiple LoRA networks beyond plain summation.
Dense modes (ties, dare_ties, dare_linear, magnitude_prune) combine the
networks' dense deltas elementwise; the result rides the normal apply tail
(side-channel hosting on sub-8-bit SDNQ, requantize at int8 and above,
direct add on unquantized layers). Select modes (klora, estlora) keep both
networks' contributions separate and choose a per-layer winner, shifting
from the first loaded network (subject) toward the second (style) across
the sampling steps. Selection scores depend only on the weights, so the
shift reduces to at most one flip per layer per generation, executed from
the step callback against a schedule finalized at apply time.
TIES arXiv:2306.01708, DARE arXiv:2311.03099, K-LoRA arXiv:2502.18461,
EST-LoRA arXiv:2508.02165 (its measured style-discrepancy estimate is
exposed as an option instead of being derived from probe generations).
"""
import time
import weakref
import hashlib
import torch
from modules import shared
from modules.logger import log
DENSE_MODES = ('ties', 'dare_ties', 'dare_linear', 'magnitude_prune')
SELECT_MODES = ('klora', 'estlora')
KLORA_BETA = 0.5 # the paper's fixed ramp offset; only the slope is user-tunable
ROW_CHUNK = 512 # fp32 interiors run in first-dim slices; also fixes the DARE draw sequence
SAMPLE_CAP = 1 << 22 # strided subsample bound for magnitude quantiles (full-size quantile exceeds torch limits)
state: dict = {'entries': {}, 'flips': {}, 'gamma': 1.0, 'gamma_e': 1.0, 'total_steps': 0, 'finalized': False, 'reported': None}
warned: set = set()
warned_context = None
def mode():
return getattr(shared.opts, 'lora_stack_mode', 'sum') or 'sum'
def density():
return float(getattr(shared.opts, 'lora_stack_density', 0.5))
def ramp_alpha():
return float(getattr(shared.opts, 'lora_stack_alpha', 0.0))
def manual_discrepancy():
return float(getattr(shared.opts, 'lora_stack_discrepancy', 0.5))
def signature():
m = mode()
if m in DENSE_MODES:
return f'{m}:{density():.2f}'
if m in SELECT_MODES:
return f'{m}:{ramp_alpha():.2f}:{manual_discrepancy():.2f}'
return 'sum'
def warn_context():
"""Settings the degradation warnings below speak about."""
return (signature(), getattr(shared.opts, 'diffusers_offload_mode', ''), int(getattr(shared.opts, 'lora_sdnq_host_rank', 0) or 0), getattr(shared.opts, 'sd_model_checkpoint', ''))
def warn_once(key, message):
global warned_context # pylint: disable=global-statement
context = warn_context()
if context != warned_context:
warned.clear() # what was reported under the old settings says nothing about the new ones
warned_context = context
if key not in warned:
warned.add(key)
log.warning(message)
def select_blocked():
return 'Model' in (getattr(shared.opts, 'cuda_compile', None) or [])
def active_dense(n_contrib):
return mode() in DENSE_MODES and n_contrib >= 2
def select_possible(n_loaded):
"""True when the loaded set could engage a select mode; silent, for the fuse gate."""
return mode() in SELECT_MODES and n_loaded == 2 and not select_blocked()
def select_engaged():
"""True while selection schedules are live on model layers."""
return bool(state['entries'])
def active_select(n_loaded):
m = mode()
if m not in SELECT_MODES:
return False
if n_loaded != 2:
log.warning(f'Network stack: mode={m} networks={n_loaded} required=2 fallback=sum')
return False
if select_blocked():
log.warning(f'Network stack: mode={m} compile=model fallback=sum')
return False
return True
def seed_for(layer_name, net_name):
payload = f'{layer_name}|{net_name}|{mode()}|{round(density(), 6)}'
return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], 'little')
def magnitude_threshold(delta, dens):
flat = delta.abs().flatten()
step = max(1, flat.numel() // SAMPLE_CAP)
return torch.quantile(flat[::step].float(), 1.0 - dens)
def dare_generator(device, layer_name, net_name):
gen = torch.Generator(device=device)
gen.manual_seed(seed_for(layer_name, net_name))
return gen
def combine(named_deltas, layer_name):
"""Combine per-network dense deltas under the active dense mode; returns a tensor in the first delta's dtype."""
m = mode()
dens = density()
deltas = [d for _, d in named_deltas]
out_dtype = deltas[0].dtype
result = torch.zeros_like(deltas[0], dtype=torch.float32)
thresholds = [magnitude_threshold(d, dens) for d in deltas] if m in ('ties', 'magnitude_prune') else [None] * len(deltas)
gens = [dare_generator(deltas[0].device, layer_name, name) for name, _ in named_deltas] if m in ('dare_ties', 'dare_linear') else [None] * len(deltas)
for start in range(0, deltas[0].shape[0], ROW_CHUNK):
stop = min(start + ROW_CHUNK, deltas[0].shape[0])
chunks = []
for i, d in enumerate(deltas):
c = d[start:stop].to(torch.float32)
if thresholds[i] is not None:
c = c * (c.abs() >= thresholds[i])
if gens[i] is not None:
keep = torch.rand(c.shape, generator=gens[i], device=c.device, dtype=torch.float32) < dens
c = c * keep / dens
chunks.append(c)
if m in ('ties', 'dare_ties'):
total = torch.stack(chunks).sum(dim=0)
elected = torch.sign(total)
agree = [c * ((torch.sign(c) == elected) & (c != 0)) for c in chunks]
count = torch.stack([(a != 0).to(torch.float32) for a in agree]).sum(dim=0).clamp(min=1.0)
result[start:stop] = torch.stack(agree).sum(dim=0) / count
else: # dare_linear, magnitude_prune: independent per-delta edits, plain sum
result[start:stop] = torch.stack(chunks).sum(dim=0)
return result.to(out_dtype)
def score_pair(d0, d1, rank0, rank1):
"""Selection scores for a dense delta pair: klora top-K sums (K = rank product) or est energies; plus abs-sums for the global balance.
Row-chunked fp32 interiors with fp64 accumulators and one device sync for
all four reductions. Full-tensor staging (fp32 copy, abs copy, top-k
workspace) peaks hundreds of MB per large layer, which collides with block
swapping on offloaded denoisers; chunking bounds the transient to the
chunk. The global top-K over per-chunk top-K candidates selects the same
element set as a whole-tensor top-K.
"""
k = max(1, int(rank0) * int(rank1)) if mode() == 'klora' else 0
accs = []
for d in (d0, d1):
score = torch.zeros((), device=d.device, dtype=torch.float64)
abs_sum = torch.zeros((), device=d.device, dtype=torch.float64)
cands = []
for start in range(0, d.shape[0], ROW_CHUNK):
c = d[start:start + ROW_CHUNK].to(torch.float32).abs() # out-of-place abs: to() may alias a caller-owned fp32 tensor
abs_sum += c.sum(dtype=torch.float64)
if k:
flat = c.flatten()
cands.append(torch.topk(flat, min(k, flat.numel()), sorted=False).values)
else:
score += c.square().sum(dtype=torch.float64)
if k and cands:
allc = torch.cat(cands) if len(cands) > 1 else cands[0]
score = torch.topk(allc, min(k, allc.numel()), sorted=False).values.sum(dtype=torch.float64)
accs.append((score, abs_sum))
packed = torch.stack([accs[0][0], accs[0][1], accs[1][0], accs[1][1]]).cpu()
return (float(packed[0]), float(packed[2])), (float(packed[1]), float(packed[3]))
def register_weight_pair(layer_name, module, per_net, wanted_names=None):
"""Score and register a weight-kind selection pair; True when the layer is scheduled.
The scores persist in the factor cache when a pass identity is given, so a
later apply of the same configuration registers from the record alone.
"""
from modules.lora import lora_common as l
if per_net is None or len(per_net) != 2:
return False
ranks, names = [], []
for net_name, d in per_net:
if d is None:
return False
net = next((n for n in l.loaded_networks if n.name == net_name), None)
net_module = net.modules.get(layer_name, None) if net is not None else None
if net_module is None:
return False
names.append(net_name)
ranks.append(int(getattr(net_module, 'dim', 0) or 0) or 64)
scores, abs_sums = score_pair(per_net[0][1], per_net[1][1], ranks[0], ranks[1])
if wanted_names is not None:
from modules.lora import lora_factor_cache
lora_factor_cache.begin_pass(wanted_names)
lora_factor_cache.store_scores(layer_name, scores, abs_sums)
register(layer_name, module, 'weight', scores, nets=tuple(names), abs_sums=abs_sums)
return True
def register_weight_pair_cached(layer_name, module, wanted_names):
"""Register a weight-kind pair from its cached score record; True when served.
The record was stored under the same configuration signature, which pins
the loaded pair, multipliers and stack settings, so both networks are known
to target the layer and the prompt-order roles are unchanged.
"""
from modules.lora import lora_common as l
from modules.lora import lora_factor_cache
if len(l.loaded_networks) != 2:
return False
lora_factor_cache.begin_pass(wanted_names)
rec = lora_factor_cache.lookup_scores(layer_name)
if rec is None:
return False
scores, abs_sums = rec
register(layer_name, module, 'weight', scores, nets=tuple(n.name for n in l.loaded_networks), abs_sums=abs_sums)
return True
def drop(layer_name):
"""Forget a layer's selection entry (its factors were removed or restored)."""
if layer_name is not None and state['entries'].pop(layer_name, None) is not None:
state['finalized'] = False
def score_energy(up, down):
"""EST layer score: squared Frobenius norm of up@down via the Gram identity, no materialization."""
u = up.to(torch.float32)
dn = down.to(torch.float32)
return float(((u.t() @ u) * (dn @ dn.t())).sum())
def clear():
state['entries'] = {}
state['flips'] = {}
state['gamma'] = 1.0
state['gamma_e'] = 1.0
state['total_steps'] = 0
state['finalized'] = False
state['reported'] = None
def register(layer_name, module, kind, scores, segments: tuple[tuple[int, int], tuple[int, int], bool] | None = None, nets=None, abs_sums=None):
"""Record a select-mode layer for schedule finalization.
kind 'factor': segments = ((s0, s1), (t0, t1), transposed) column ranges on the svd
channel; both segments' pristine values are stashed for flips. kind 'weight': nets =
the two network names; the winner delta is recomputed from the layer backup at
selection time. abs_sums feeds the global magnitude balance (klora gamma).
"""
entry = {'layer': layer_name, 'module': weakref.ref(module), 'kind': kind, 'segments': segments, 'scores': scores, 'nets': nets, 'abs_sums': abs_sums, 'stash': None}
if kind == 'factor':
if segments is None:
raise ValueError("segments is required when kind='factor'")
(s0, s1), (t0, t1), transposed = segments
up = module.svd_up.data
entry['stash'] = (segment_view(up, s0, s1, transposed).clone(), segment_view(up, t0, t1, transposed).clone())
state['entries'][layer_name] = entry
state['finalized'] = False
def segment_view(up, start, stop, transposed):
return up[start:stop] if transposed else up[:, start:stop]
def layer_flip_step(scores, total_steps):
"""First step index at which the style side wins; total_steps when it never does, 0 when style wins from the start."""
m = mode()
sc, ss = scores
for step in range(total_steps):
t = step / max(1, total_steps - 1)
if m == 'klora':
ramp = state['gamma'] * (ramp_alpha() * t + KLORA_BETA)
if ss * ramp > sc:
return step
else: # estlora: content keeps the layer while sc >= gamma_t * ss
# est energies are ||dW||^2, so a magnitude gap enters squared; balance the style side by
# the total-energy ratio (mirrors klora's gamma) so the louder adapter cannot win on scale alone
ramp = ramp_alpha() * t + (1.0 - manual_discrepancy())
if sc < ramp * ss * state['gamma_e']:
return step
return total_steps
def materialize_model():
"""Weight-kind selection rewrites module weights outside the activation walk; rebuild balanced-offload modules real first (mirrors network_activate)."""
from modules import sd_models
if getattr(shared.opts, 'diffusers_offload_mode', None) == 'balanced' and getattr(shared, 'sd_model', None) is not None:
sd_models.apply_balanced_offload(shared.sd_model, force=True, silent=True)
def finalize(total_steps):
"""Build the inverted flip map for the pass; select-mode layers start at their step-0 winner."""
state['total_steps'] = int(total_steps)
# both balances derive from the live entries every time, so drops and re-registrations stay consistent by construction
num = sum(e['abs_sums'][0] for e in state['entries'].values() if e['abs_sums'] is not None)
den = sum(e['abs_sums'][1] for e in state['entries'].values() if e['abs_sums'] is not None)
state['gamma'] = (num / den) if den > 0 else 1.0
e_num = sum(e['scores'][0] for e in state['entries'].values()) # est scores ARE the per-layer energies; their totals give the scale-invariant balance
e_den = sum(e['scores'][1] for e in state['entries'].values())
state['gamma_e'] = (e_num / e_den) if e_den > 0 else 1.0
state['flips'] = {}
stats = {'weight_n': 0, 'factor_n': 0, 'materialize': 0.0, 'select': 0.0, 'w_move': 0.0, 'w_calc': 0.0, 'w_apply': 0.0}
state['stats'] = stats
stats['weight_n'] = sum(1 for e in state['entries'].values() if e['kind'] == 'weight')
stats['factor_n'] = len(state['entries']) - stats['weight_n']
if stats['weight_n'] > 0:
t0 = time.time()
materialize_model()
stats['materialize'] = time.time() - t0
style_first = 0
t0 = time.time()
for layer_name, entry in list(state['entries'].items()): # snapshot: apply_selection drops entries whose module died
flip_at = layer_flip_step(entry['scores'], state['total_steps'])
initial = 1 if flip_at == 0 else 0
style_first += initial
apply_selection(layer_name, entry, initial)
if 0 < flip_at < state['total_steps']:
state['flips'].setdefault(flip_at - 1, []).append(layer_name) # step callbacks fire after the denoise, so the flip runs one step early to be live during the crossover step's forward
stats['select'] = time.time() - t0
state['finalized'] = True
if len(state['entries']) > 0: # only a built schedule can carry a flip count, so this is the line that shows selection is live rather than requested
gamma = state['gamma_e'] if mode() == 'estlora' else state['gamma']
report = (mode(), len(state['entries']), style_first, sum(len(v) for v in state['flips'].values()), state['total_steps'], round(gamma, 3))
if report != state['reported']: # rebuilt every pass, so a batch would otherwise repeat one line per image
state['reported'] = report
log.info(f'Network load: type=LoRA stack={report[0]} layers={report[1]} style={report[2]} flips={report[3]} steps={report[4]} gamma={report[5]:.3f}')
# logged every pass: the reset runs outside the activate walk, so its cost is invisible to the load timers
log.debug(f'Network select: type=LoRA reset weight={stats["weight_n"]} factor={stats["factor_n"]} time={{materialize: {stats["materialize"]:.2f}, select: {stats["select"]:.2f}, move: {stats["w_move"]:.2f}, calc: {stats["w_calc"]:.2f}, apply: {stats["w_apply"]:.2f}}}')
def reset(total_steps):
"""Per-pass reset from set_callbacks_p: restore initial selections and reschedule for this pass's step count."""
if mode() not in SELECT_MODES or not state['entries'] or int(total_steps) <= 0:
return
finalize(total_steps)
def on_step(step):
"""Flip the layers whose crossover is this step; non-flip steps are a dict miss."""
if not state['finalized']:
return
layers = state['flips'].get(int(step), ())
if not layers:
return
t0 = time.time()
for layer_name in layers:
entry = state['entries'].get(layer_name)
if entry is not None:
apply_selection(layer_name, entry, 1)
log.debug(f'Network select: type=LoRA flip step={int(step)} layers={len(layers)} time={time.time() - t0:.2f}')
def apply_selection(layer_name, entry, winner):
module = entry['module']()
if module is None:
state['entries'].pop(layer_name, None)
return
if entry['kind'] == 'factor':
(s0, s1), (t0, t1), transposed = entry['segments']
up = module.svd_up.data
keep_seg, drop_seg = ((t0, t1), (s0, s1)) if winner == 1 else ((s0, s1), (t0, t1))
stash = entry['stash'][winner]
segment_view(up, keep_seg[0], keep_seg[1], transposed).copy_(stash.to(device=up.device, dtype=up.dtype))
segment_view(up, drop_seg[0], drop_seg[1], transposed).zero_()
else:
weight_selection(module, entry, winner)
def weight_selection(module, entry, winner):
from modules.lora import lora_common as l
from modules.lora.lora_apply import network_apply_weights
if getattr(module, 'sdnq_dequantizer', None) is not None:
warn_once('select-sdnq-weight', 'Network stack: flip=skipped layer=quantized') # quantized backups are packed tensors; only the segment path can flip them
return
backup = getattr(module, 'network_weights_backup', None)
if not isinstance(backup, torch.Tensor): # fuse mode keeps a bool sentinel, not a pristine copy
warn_once('select-nobackup', 'Network stack: flip=skipped backup=none')
return
net = next((n for n in l.loaded_networks if n.name == entry['nets'][winner]), None)
net_module = net.modules.get(entry['layer'], None) if net is not None else None
if net_module is None:
return
weight = getattr(module, 'weight', None)
if weight is None or weight.is_meta:
warn_once('select-offloaded', 'Network stack: flip=skipped weight=offloaded')
return
from modules import devices
stats = state.get('stats') or {}
device = weight.device
t0 = time.time()
base = backup.to(devices.device) # a swapped-out layer keeps its weight on cpu; the delta matmul belongs on the accelerator regardless
t1 = time.time()
updown = net_module.calc_updown(base)[0].to(device)
t2 = time.time()
network_apply_weights(module, updown, None, device=device) # recomputes from the pristine backup, requantizing where the layer needs it
stats['w_move'] = stats.get('w_move', 0.0) + (t1 - t0)
stats['w_calc'] = stats.get('w_calc', 0.0) + (t2 - t1)
stats['w_apply'] = stats.get('w_apply', 0.0) + (time.time() - t2)
modules/lora/lora_timers.py
-5
View File
@@ -5,7 +5,6 @@ class Timer:
calc: float = 0
apply: float = 0
move: float = 0
restore: float = 0
activate: float = 0
deactivate: float = 0
@@ -26,13 +25,9 @@ class Timer:
self.calc = 0
self.apply = 0
self.move = 0
self.restore = 0
if complete:
self.activate = 0
self.deactivate = 0
def add(self, name, t):
self.__dict__[name] += t
def __str__(self):
return f'{self.__class__.__name__}({self.summary})'
modules/lora/native_adapter.py
+119 -101
View File
@@ -22,9 +22,8 @@ of a fused weight is described) and the per-arch ``resolve_targets`` callable
each loader passes in (how a parsed ``(prefix, base)`` maps to one or more
diffusers paths plus optional chunk descriptors).
Per-arch loader modules import this module and pass their own ``prefixes``,
``bare_prefixes``, ``bare_diffusers_prefixes``, and ``resolve_targets`` to the
generic helpers.
Per-arch loader modules import this module and pass their own ``prefixes``
and ``resolve_targets`` to the generic helpers.
"""
import os
@@ -33,7 +32,7 @@ from dataclasses import dataclass
import torch
from modules import shared, sd_models, sd_models_utils
from modules.sd_models import read_state_dict # pylint: disable=unused-import
from modules.logger import log
from modules.lora import (
lora_convert, network, network_boft, network_full, network_glora,
@@ -55,10 +54,11 @@ from modules.lora import lora_common as l
KNOWN_PREFIXES_DEFAULT = ("diffusion_model.", "transformer.", "lora_unet_", "lora_transformer_", "lycoris_")
# Sentinel ``prefix_used`` value emitted by :func:`parse_key` when a bare path
# starting with a member of ``bare_diffusers_prefixes`` matches. A loader
# ``resolve_targets`` may dispatch on this string to rewrite the base path;
# when it declines, :func:`resolve_group_targets` binds the path verbatim.
# Sentinel ``prefix_used`` value emitted by :func:`parse_key` for a bare path,
# one that matched no arch prefix. A loader ``resolve_targets`` may dispatch on
# this string to rewrite the base path; when it declines,
# :func:`resolve_group_targets` binds the path verbatim, and a path naming no
# live module counts as unmapped instead of vanishing.
BARE_DIFFUSERS_PREFIX_USED = "bare_diffusers"
@@ -92,6 +92,11 @@ def _resolve_prefix(network_prefix, prefix_used):
SUFFIX_NORMALIZE = {
"lora_A.weight": "lora_down.weight",
"lora_B.weight": "lora_up.weight",
# bare parameter names, saved by wrappers that hold the factors as nn.Parameter (alibaba-pai PDD files)
"lora_down": "lora_down.weight",
"lora_a": "lora_down.weight",
"lora_b": "lora_up.weight",
"lora_up": "lora_up.weight",
}
@@ -103,6 +108,8 @@ SUFFIX_NORMALIZE = {
LORA_SUFFIXES = (
".lora_down.weight", ".lora_up.weight", ".lora_mid.weight",
".lora_A.weight", ".lora_B.weight",
".lora_down", ".lora_up",
".lora_a", ".lora_b", # lowercase peft factor names without .weight (TaoLive adapters)
# diff_b: bias delta some saves pair with the weight LoRA, applied as ex_bias.
# magnitude / lora_magnitude_vector: DoRA row norms (ai-toolkit / PEFT key
# names); converted onto the dora_scale path by try_load_lora.
@@ -154,7 +161,7 @@ FULL_SUFFIXES = (
# on accidental overlaps with other families.
LORA_MARKERS = (
".lora_down.weight", ".lora_up.weight",
".lora_down", ".lora_up", ".lora_a", ".lora_b", # bare and .weight forms alike
".lora_A.weight", ".lora_B.weight",
# PEFT named-adapter saves embed the slot name as ``.lora_A.<name>.weight``;
# the trailing-dot forms catch every variant.
@@ -174,9 +181,9 @@ FULL_MARKERS = (".diff",)
@dataclass(frozen=True)
class ChunkSpec:
"""How to slice a fused weight along dim 0 for one target module.
"""How to take a fused weight's rows along dim 0 for one target module.
Two forms supported:
Three forms; the reorder composes with either slice:
- Equal chunks (``idx`` + ``total``): fused QKV split into Q/K/V via
``torch.chunk(up, total, dim=0)[idx]``. Used by flux2 / z-image where
@@ -185,6 +192,11 @@ class ChunkSpec:
``up[start:end]``. Used by chroma's single-block ``linear1`` which
fuses Q / K / V / proj_mlp at unequal sizes
(``[3072, 3072, 3072, 12288]``).
- Row reorder (``reorder``): the module lays out equal row blocks in a
different order from the save; ``(1, 0)`` swaps the halves of a fused
SwiGLU projection saved ``[gate; value]`` onto a ``[value; gate]``
module. Applied to the rows the slice selects. Only the LoRA family
permutes rows; the others skip a reordered target.
Generic loaders check :attr:`is_equal_chunks` to decide between the two
forms and select the appropriate ``NetworkModule*Chunk`` /
@@ -194,11 +206,16 @@ class ChunkSpec:
total: int | None = None
start: int | None = None
end: int | None = None
reorder: tuple[int, ...] | None = None
@property
def is_equal_chunks(self) -> bool:
return self.idx is not None and self.total is not None
@property
def is_slice(self) -> bool:
return self.is_equal_chunks or self.start is not None
# === Key normalizations (applied universally by parse_key) ===
@@ -248,6 +265,7 @@ def has_marker(state_dict, markers):
def resolve_mapping():
from modules import shared
"""Ensure ``network_layer_mapping`` is populated, return it (or empty dict)."""
sd_model = getattr(shared.sd_model, "pipe", shared.sd_model)
lora_convert.assign_network_names_to_compvis_modules(sd_model)
@@ -288,19 +306,19 @@ def finalize_network(net, name, family, lora_scale, t0, unmapped=0, mismatch=0,
return net
def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
"""LoRA-style rank-and-dim sanity check against the live module weight.
Honors SDNQ-quantized modules by reading the original shape from the
dequantizer rather than the packed weight tensor.
"""
def module_shape(sd_module):
"""The live weight shape of a module, read from the dequantizer for SDNQ-quantized layers; None without a weight."""
if not hasattr(sd_module, "weight"):
return False
return None
if hasattr(sd_module, "sdnq_dequantizer"):
mod_shape = sd_module.sdnq_dequantizer.original_shape
else:
mod_shape = sd_module.weight.shape
if len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2:
return tuple(sd_module.sdnq_dequantizer.original_shape)
return tuple(sd_module.weight.shape)
def shapes_match(sd_module, down_w: torch.Tensor, up_w: torch.Tensor) -> bool:
"""LoRA-style rank-and-dim sanity check against the live module weight."""
mod_shape = module_shape(sd_module)
if mod_shape is None or len(mod_shape) < 2 or len(down_w.shape) < 2 or len(up_w.shape) < 2:
return False
return down_w.shape[1] == mod_shape[1] and up_w.shape[0] == mod_shape[0]
@@ -368,7 +386,7 @@ def lokr_shapes_match(sd_module, kron_shape, chunk: ChunkSpec | None) -> bool:
kron_out, kron_in_flat = kron_shape
if kron_in_flat != mod_in_flat:
return False
if chunk is None:
if chunk is None or not chunk.is_slice:
return kron_out == mod_shape[0]
if chunk.is_equal_chunks:
return kron_out == mod_shape[0] * chunk.total
@@ -378,15 +396,15 @@ def lokr_shapes_match(sd_module, kron_shape, chunk: ChunkSpec | None) -> bool:
# === Parsing primitives ===
def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()):
def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT):
"""Return ``(prefix_used, base, suffix_normalized)`` or ``None``.
``prefix_used`` is the matched element of ``prefixes``, ``BARE_DIFFUSERS_PREFIX_USED``
if a member of ``bare_diffusers_prefixes`` matched, or ``None`` for a key
that matched a member of ``bare_prefixes``. ``base`` is the path with prefix
and suffix removed. ``suffix_normalized`` is the suffix (without the leading
dot) after applying :data:`SUFFIX_NORMALIZE` (e.g. ``lora_A.weight`` becomes
``lora_down.weight``).
``prefix_used`` is the matched element of ``prefixes``, or
``BARE_DIFFUSERS_PREFIX_USED`` for a bare key, which the loader offers to
the resolver and counts as unmapped when nothing binds. ``base`` is the
path with prefix and suffix removed. ``suffix_normalized`` is the suffix
(without the leading dot) after applying :data:`SUFFIX_NORMALIZE` (e.g.
``lora_A.weight`` becomes ``lora_down.weight``).
Always applies :func:`unwrap_peft_wrapper` and :func:`strip_peft_adapter_name`
to the raw key before format detection so callers do not have to opt in.
@@ -401,10 +419,7 @@ def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(
stripped = key[len(p):]
break
if prefix_used is None:
if any(key.startswith(p) for p in bare_diffusers_prefixes):
prefix_used = BARE_DIFFUSERS_PREFIX_USED
elif not any(key.startswith(p) for p in bare_prefixes):
return None
prefix_used = BARE_DIFFUSERS_PREFIX_USED
matched_suffix = None
split_at = -1
@@ -424,7 +439,7 @@ def parse_key(key, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(
return prefix_used, base, suffix
def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT, bare_prefixes=(), bare_diffusers_prefixes=()):
def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT):
"""Group state-dict entries by ``(prefix_used, base)``.
Returns ``{(prefix_used, base): {suffix: tensor, ...}}`` where each suffix
@@ -434,12 +449,7 @@ def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT,
"""
groups: dict[tuple, dict[str, torch.Tensor]] = {}
for key, value in state_dict.items():
parsed = parse_key(
key, suffixes,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
parsed = parse_key(key, suffixes, prefixes=prefixes)
if parsed is None:
continue
prefix_used, base, suffix = parsed
@@ -451,11 +461,6 @@ def group_by_suffixes(state_dict, suffixes, *, prefixes=KNOWN_PREFIXES_DEFAULT,
return groups
# Surface ``sd_models.read_state_dict`` here so loader modules don't have to
# import ``sd_models`` directly; keeps the per-arch wrapper imports compact.
read_state_dict = sd_models.read_state_dict
def resolve_group_targets(resolve_targets, prefix_used, base):
"""Map a parsed ``(prefix_used, base)`` group to ``[(diffusers_path, chunk), ...]``.
@@ -508,20 +513,27 @@ def resolve_group_targets(resolve_targets, prefix_used, base):
def slice_chunk_rows(t, chunk: ChunkSpec):
"""Slice dim 0 of ``t`` per ``chunk``.
"""Slice dim 0 of ``t`` per ``chunk``, then lay the selected rows out in the chunk's order.
Equal-chunks form uses ``torch.chunk`` (faster for the symmetric case);
row-range form uses tensor slicing for arbitrary partitions.
"""
if chunk.is_equal_chunks:
return torch.chunk(t, chunk.total, dim=0)[chunk.idx].contiguous()
return t[chunk.start:chunk.end].contiguous()
t = torch.chunk(t, chunk.total, dim=0)[chunk.idx]
elif chunk.start is not None:
t = t[chunk.start:chunk.end]
if chunk.reorder is not None:
blocks = torch.chunk(t, len(chunk.reorder), dim=0)
t = torch.cat([blocks[i] for i in chunk.reorder], dim=0)
return t.contiguous()
def _slice_lora_chunk(w, chunk: ChunkSpec):
"""Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``."""
def slice_lora_chunk(w, chunk: ChunkSpec):
"""Return a shallow copy of ``w`` with ``lora_up.weight`` sliced per ``chunk``; a dense bias follows a pure reorder."""
out = dict(w)
out["lora_up.weight"] = slice_chunk_rows(w["lora_up.weight"], chunk)
if "bias" in w and not chunk.is_slice:
out["bias"] = slice_chunk_rows(w["bias"], chunk)
return out
@@ -564,13 +576,22 @@ def slice_bias_delta(w, chunk: ChunkSpec, fused_out):
def try_load_lora(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
network_alpha=None,
adapt_weights=None,
arch_name="generic"):
"""Generic LoRA loader (handles DoRA via the universal ``finalize_updown`` hook).
Fused targets are chunked at load time by slicing ``lora_up`` along dim 0;
the down-side is shared across the resolved targets.
``network_alpha`` is a file-level alpha for files without alpha tensors;
a file carrying any alpha of its own keeps those and ignores it.
``adapt_weights(sd_module, network_key, w)`` lets an arch refit a delta onto
a module whose live layout differs from the trained one (a pruned AdaLN
basis, for instance) before the shape check; returning None keeps ``w``.
"""
t0 = time.time()
state_dict = read_state_dict(network_on_disk.filename, what="network")
@@ -579,12 +600,12 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, LORA_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
if network_alpha is not None and any("alpha" in w for w in groups.values()):
network_alpha = None
unmapped = 0
mismatch = 0
@@ -592,6 +613,9 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
for (prefix, base), w in groups.items():
if "lora_down.weight" not in w or "lora_up.weight" not in w:
continue
if network_alpha is not None:
w = dict(w)
w["alpha"] = torch.tensor(float(network_alpha))
# DoRA magnitude vectors: ai-toolkit saves `magnitude`, PEFT/diffusers
# `lora_magnitude_vector`. Both are 1-D per-output row norms with
# dora_scale semantics; reshape to (out, 1) so the apply-time
@@ -612,14 +636,14 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
target_w = w
if chunk is not None:
if "bias" in w or "bias_indices" in w:
if "bias_indices" in w or ("bias" in w and chunk.is_slice):
# Weight-shaped bias residuals (dense or LyCORIS sparse
# triplet) are not partitioned onto fused targets.
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
continue
fused_out = w["lora_up.weight"].shape[0]
target_w = _slice_lora_chunk(w, chunk)
target_w = slice_lora_chunk(w, chunk)
target_w = slice_dora_scale(target_w, chunk, fused_out)
if target_w is None:
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} per-input DoRA on fused target skipped (unsupported)')
@@ -631,22 +655,20 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
skipped += 1
continue
if adapt_weights is not None:
target_w = adapt_weights(sd_module, network_key, target_w) or target_w
if not shapes_match(sd_module, target_w["lora_down.weight"], target_w["lora_up.weight"]):
log.warning(
f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key}'
f' lora={target_w["lora_down.weight"].shape[1]}x{target_w["lora_up.weight"].shape[0]}'
f' module={getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}'
f' shape mismatch'
)
if l.debug:
_module = f'{getattr(sd_module, "weight", None).shape if hasattr(sd_module, "weight") else "?"}'
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} lora={target_w["lora_down.weight"].shape[1]}x{target_w["lora_up.weight"].shape[0]} module={_module} shape mismatch')
mismatch += 1
continue
if "diff_b" in target_w and not bias_delta_fits(sd_module, target_w["diff_b"]):
log.warning(
f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key}'
f' bias={tuple(target_w["diff_b"].shape)} module={tuple(sd_module.bias.shape)}'
f' bias shape mismatch'
)
if l.debug:
_bias = f'bias={tuple(target_w["diff_b"].shape)} module={tuple(sd_module.bias.shape)}'
log.warning(f'Network load: type=LoRA name="{name}" arch={arch_name} key={network_key} {_bias} bias shape mismatch')
mismatch += 1
continue
@@ -658,8 +680,8 @@ def try_load_lora(name, network_on_disk, lora_scale, *,
def try_load_lokr(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic LoKR loader.
@@ -677,11 +699,9 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, LOKR_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -711,6 +731,10 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
continue
target_w = w
if chunk is not None:
if chunk.reorder is not None:
log.warning(f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key} row reorder on fused target skipped (unsupported)')
skipped += 1
continue
if "bias" in w:
log.warning(f'Network load: type=LoKR name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
@@ -739,8 +763,8 @@ def try_load_lokr(name, network_on_disk, lora_scale, *,
def try_load_loha(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic LoHA (Hadamard product) loader.
@@ -757,11 +781,9 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, LOHA_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -773,7 +795,7 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
targets = resolve_group_targets(resolve_targets, prefix, base)
is_fused = any(t[1] is not None for t in targets)
if is_fused and is_tucker:
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={base} Tucker fused QKV skipped (unsupported)')
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={base} Tucker fused target skipped (unsupported)')
skipped += 1
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
@@ -785,6 +807,10 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
continue
target_w = w
if chunk is not None:
if chunk.reorder is not None:
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} row reorder on fused target skipped (unsupported)')
skipped += 1
continue
if "bias" in w:
log.warning(f'Network load: type=LoHA name="{name}" arch={arch_name} key={network_key} weight-shaped bias on fused target skipped (unsupported)')
skipped += 1
@@ -808,8 +834,8 @@ def try_load_loha(name, network_on_disk, lora_scale, *,
def try_load_oft(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic OFT/BOFT loader.
@@ -832,11 +858,9 @@ def try_load_oft(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, OFT_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -847,7 +871,7 @@ def try_load_oft(name, network_on_disk, lora_scale, *,
is_boft = "oft_blocks" in w and w["oft_blocks"].ndim == 4
targets = resolve_group_targets(resolve_targets, prefix, base)
if any(t[1] is not None for t in targets):
log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
log.warning(f'Network load: type={"BOFT" if is_boft else "OFT"} name="{name}" arch={arch_name} key={base} fused target skipped (unsupported)')
skipped += 1
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
@@ -868,8 +892,8 @@ def try_load_oft(name, network_on_disk, lora_scale, *,
def try_load_ia3(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic IA3 loader.
@@ -891,11 +915,9 @@ def try_load_ia3(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, IA3_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -905,7 +927,7 @@ def try_load_ia3(name, network_on_disk, lora_scale, *,
continue
targets = resolve_group_targets(resolve_targets, prefix, base)
if any(t[1] is not None for t in targets):
log.warning(f'Network load: type=IA3 name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
log.warning(f'Network load: type=IA3 name="{name}" arch={arch_name} key={base} fused target skipped (unsupported)')
skipped += 1
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
@@ -923,8 +945,8 @@ def try_load_ia3(name, network_on_disk, lora_scale, *,
def try_load_glora(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic GLoRA loader.
@@ -942,11 +964,9 @@ def try_load_glora(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, GLORA_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -956,7 +976,7 @@ def try_load_glora(name, network_on_disk, lora_scale, *,
continue
targets = resolve_group_targets(resolve_targets, prefix, base)
if any(t[1] is not None for t in targets):
log.warning(f'Network load: type=GLoRA name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
log.warning(f'Network load: type=GLoRA name="{name}" arch={arch_name} key={base} fused target skipped (unsupported)')
skipped += 1
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
@@ -974,8 +994,8 @@ def try_load_glora(name, network_on_disk, lora_scale, *,
def try_load_norm(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"): # pylint: disable=unused-argument
"""Generic Norm (LayerNorm / RMSNorm weight + bias delta) loader.
@@ -996,11 +1016,9 @@ def try_load_norm(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, NORM_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -1039,8 +1057,8 @@ def try_load_norm(name, network_on_disk, lora_scale, *,
def try_load_full(name, network_on_disk, lora_scale, *,
resolve_targets, prefixes=KNOWN_PREFIXES_DEFAULT,
bare_prefixes=(), bare_diffusers_prefixes=(),
network_prefix=NETWORK_PREFIX_DEFAULT,
group_by_suffixes_fn=group_by_suffixes,
arch_name="generic"):
"""Generic Full (full-rank weight delta) loader.
@@ -1057,11 +1075,9 @@ def try_load_full(name, network_on_disk, lora_scale, *,
mapping = resolve_mapping()
net = new_network(name, network_on_disk)
groups = group_by_suffixes(
groups = group_by_suffixes_fn(
state_dict, FULL_SUFFIXES,
prefixes=prefixes,
bare_prefixes=bare_prefixes,
bare_diffusers_prefixes=bare_diffusers_prefixes,
)
unmapped = 0
@@ -1072,7 +1088,7 @@ def try_load_full(name, network_on_disk, lora_scale, *,
continue
targets = resolve_group_targets(resolve_targets, prefix, base)
if any(t[1] is not None for t in targets):
log.warning(f'Network load: type=Full name="{name}" arch={arch_name} key={base} fused QKV skipped (unsupported)')
log.warning(f'Network load: type=Full name="{name}" arch={arch_name} key={base} fused target skipped (unsupported)')
skipped += 1
continue
arch_prefix = _resolve_prefix(network_prefix, prefix)
@@ -1114,6 +1130,7 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders):
tuple of partial-applied generic loaders, each already bound to the arch's
``resolve_targets`` and prefix tuples.
"""
from modules import sd_models_utils
sd_models_utils.state_dict_cache.enable()
net = None
mismatch = 0
@@ -1126,6 +1143,7 @@ def try_load_chain(name, network_on_disk, lora_scale, family_loaders):
net = sub
else:
net.modules.update(sub.modules)
net.extras.update(sub.extras)
sd_models_utils.state_dict_cache.disable()
if net is not None and mismatch > 0: # applying only the layers that fit leaves the model in a state nothing was trained for
log.error(f'Network load: type=LoRA name="{name}" modules={len(net.modules)} mismatch={mismatch} shapes do not match the loaded model')
modules/lora/network.py
+29 -13
View File
@@ -22,11 +22,11 @@ class SdVersion(enum.Enum):
class NetworkOnDisk:
def __init__(self, name, filename):
def __init__(self, name: str, filename: str):
self.shorthash = None
self.hash = None
self.name = name
self.filename = filename
self.name: str = name
self.filename: str = filename
if filename.startswith(shared.cmd_opts.lora_dir):
# strip("/") missed Windows's leading backslash after the slice; normalize separators
# so the registry key is one canonical form on every OS.
@@ -76,6 +76,10 @@ class NetworkOnDisk:
return 'anima'
if base.startswith('qwen'):
return 'qwen'
if base.startswith('krea2'):
return 'krea2'
if base.startswith('minimax'):
return 'minimax'
if arch.startswith("stable-diffusion-v1"):
return 'sd1'
@@ -83,7 +87,7 @@ class NetworkOnDisk:
return 'xl'
if arch.startswith("stable-cascade"):
return 'sc'
if arch.startswith("flux2") or "klein" in arch:
if arch.startswith("flux2") or arch.startswith("flux-2") or ("klein" in arch):
return 'f2'
if arch.startswith("flux"):
return 'f1'
@@ -91,12 +95,18 @@ class NetworkOnDisk:
return 'hv'
if arch.startswith("chroma"):
return 'chroma'
if arch.startswith('wan'):
return 'wan'
if arch.startswith('anima'):
return 'anima'
if arch.startswith('krea2'):
return 'krea2'
if "v1-5" in str(self.metadata.get('ss_sd_model_name', "")):
return 'sd1'
if str(self.metadata.get('ss_v2', "")) == "True":
return 'sd2'
if 'klein' in self.name.lower() or 'klein' in self.fullname.lower():
if 'klein' in self.name.lower() or ('klein' in self.fullname.lower()):
return 'f2'
if 'flux' in self.name.lower():
return 'f1'
@@ -148,8 +158,10 @@ class Network: # LoraModule
self.te_multiplier = 1.0
self.unet_multiplier = [1.0] * 3
self.dyn_dim = None
self.block_spec = None # raw lbw= value; per-layer factors resolve through lora_blocks
self.pending_config = None # staged multipliers; network_activate promotes them after the removal pass so fuse removal subtracts the delta that was applied
self.modules = {}
self.extras = {} # non-delta payloads a family carries, e.g. parallel heads
self.mismatch = 0 # deltas dropped for not fitting their target module; try_load_chain refuses the file when non-zero
self.bundle_embeddings = {}
self.mtime = None
@@ -195,15 +207,19 @@ class NetworkModule:
def multiplier(self):
unet_multiplier = 3 * [self.network.unet_multiplier] if not isinstance(self.network.unet_multiplier, list) else self.network.unet_multiplier
if self.sd_key.startswith('lora_te') or 'transformer' in self.sd_key[:20]:
return self.network.te_multiplier
if "down_blocks" in self.sd_key:
return unet_multiplier[0]
if "mid_block" in self.sd_key:
return unet_multiplier[1]
if "up_blocks" in self.sd_key:
return unet_multiplier[2]
base = self.network.te_multiplier
elif "down_blocks" in self.sd_key:
base = unet_multiplier[0]
elif "mid_block" in self.sd_key:
base = unet_multiplier[1]
elif "up_blocks" in self.sd_key:
base = unet_multiplier[2]
else:
return unet_multiplier[0]
base = unet_multiplier[0]
if getattr(self.network, 'block_spec', None) is None: # per-block strength is off for this network; no shared access on this path
return base
from modules.lora import lora_blocks
return base * lora_blocks.factor(self.sd_key, self.network)
def calc_scale(self):
if self.scale is not None:

Some files were not shown because too many files have changed in this diff Show More