diff --git a/CHANGELOG.md b/CHANGELOG.md
index 9f2dca43d..537a566ee 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,66 @@
# Change Log for SD.Next
+## Update for 2026-08-07
+
+### Highlights for 2026-08-07
+
+This release brings **Sefi-Image** and **Mage-Flow** models, plus a new **Nunchaku-Lite** inference engine
+*What else*?
+- On the server side, there are quite a few *under-the-hood* improvements, including optimized startup, optimized webserver, end-to-end profiling, storage analyzer, etc.
+- There are also several new auxiliary models, such as **Lucida** for background removal
+- And video processing now supports scripts such as prompt enhance, nudenet, etc.
+- Plus several quality-of-life improvements (better progress monitoring for one) and bug-fixes across the board
+- Updated [SD.Next Launcher](https://github.com/vladmandic/sdnext-launcher/releases/tag/v0.1.6) with improved platform compatibility and upgrade workflows
+
+*Note*: This release follows previous minor service-release which did not get full announcement, so if you missed it, check it out
+
+[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-08-07
+
+- **Models**
+ - [SeFi-Image](https://huggingface.co/SeFi-Image/SeFi-Image-5B-RL) in *Base*, *Turbo* (distilled) and *RL* (finetuned) variants
+ SeFi is an interesting model that separates generation into semantic and texture latent streams
+ and denoising semantic structure slightly ahead of texture details
+ SeFi comes in sizes with 1B, 2B and 5B params
+ - [Microsoft Mage-Flow](https://huggingface.co/mage-flow-community/Mage-Flow) in *Base* and *Turbo* (distilled) variants
+ Mage-Flow is a 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing
+ *note*: Microsoft released and then unpublished the model, but we still have a mirror available for download
+ - [Nunchaku-Lite](https://huggingface.co/lite-infer) pre-quantized models
+ included: *Z-Image, Flux.1-Dev/Schnell/Krea/Kontex, Qwen-Image/Image-Edit, Ernie-Image*
+- **Features**
+ - [Nunchaku-Lite](https://github.com/rootonchair/nunchaku-lite) inference engine
+ unlike Nunchaku, Nunchaku-Lite is based on Kernels and does not require any additional packages to be installed
+ but like original Nunchaku, it is only available for CUDA and right now only for `torch==2.11/2.12`
+ - [Krea2](https://huggingface.co/krea/Krea-2-Turbo) add *Inpaint* pipeline
+ this also makes K2 compatible with *Detailer* workflow
+ - storage analyzer: new feature that analyzes your storage used by sdnext per type and location
+ *system -> storage*
+ - video: support for scripts/extensions
+ video processing now supports scripts and extensions (if they support video processing)
+ *example*: use nudenet to automatically censor video frames :)
+ - prompt enhance: support for video generation
+ - startup: optimized server startup
+ - process: preserve audio when processing video
+ - separate progress reporting and live-preview for much more precise progress reporting
+ - add progress details to performance status bar (below the preview image)
+ - remove background: new [lucida](https://huggingface.co/egeorcun/lucida) model
+ - profile flag now logs all http requests and internal tasks
+- **API**
+ - add `/sdapi/v1/storage` endpoint to return storage usage info
+- **Internal**
+ - switch internal server to explicit `uvicorn`
+ - update core requirements
+- **Fixes**
+ - seedvr quality
+ - skip-all do not skip env init
+ - sdnq check contiguous
+ - torch reset compile cache on reload
+ - bypass sdna for caption/prompt-enhance calls
+ - skip sdnq for small weights
+ - server monitor keep websocket open
+ - unauthenticated path traversal in /thumbs
+
## Update for 2026-07-23
Primarily a service release with updates to compute packages: torch, CUDA, ROCm, etc.
diff --git a/TODO.md b/TODO.md
index 74e36f08d..78b2c9ffb 100644
--- a/TODO.md
+++ b/TODO.md
@@ -12,17 +12,14 @@
- Control tab verify overrides handling, @vladmandic
- Cloud providers, @CalamitousFelicitousness
- Video processing add/verify full API support, @CalamitousFelicitousness
-- Storage analyzer, @vladmandic
- Lora: new handler, @CalamitousFelicitousness
-- Vide: full prompt enhance
- Processing -> Video capabilities, @vladmandic
- - `NudeNet` in processing
- `RIFE` in processing
### Unassigned
-- [Nunchaku Lite](https://github.com/huggingface/diffusers/pull/14100)
- [Object clear](https://huggingface.co/jixin0101/ObjectClear) remover for Kanvas
+- [MiniMax H3](https://github.com/huggingface/diffusers/pull/14355)
- Video models: add to Reference
- Video models: support custom entries, finetunes
- UI Lite vs Expert mode
diff --git a/cli/hf-info.py b/cli/hf-info.py
index d9dfd8ebd..0da4aae1f 100755
--- a/cli/hf-info.py
+++ b/cli/hf-info.py
@@ -702,6 +702,7 @@ def search(repo_id: str) -> int:
"pipeline": pipeline_value,
"gated": gated_value,
"size": size_total_raw,
+ "size_gb": round(size_total_raw / (1024**3), 2) if isinstance(size_total_raw, int) else None,
"class": model_class,
"dit": ", ".join(main_dit_entries) if len(main_dit_entries) > 0 else None,
"dit_params": model_params_raw,
diff --git a/cli/sdnq-attention-benchmark.py b/cli/sdnq-attention-benchmark.py
index 01cc51db8..e774d0344 100755
--- a/cli/sdnq-attention-benchmark.py
+++ b/cli/sdnq-attention-benchmark.py
@@ -50,8 +50,9 @@ import signal
import logging
import argparse
import tempfile
+import statistics
+import importlib
import importlib.metadata
-import importlib.import_module
from contextlib import contextmanager
import torch
@@ -660,8 +661,16 @@ 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
-# per-test z holding the family-wise confidence at 90% across k tested candidates (sidak)
-sidak_z = {1: 1.28, 2: 1.63, 3: 1.82, 4: 1.95}
+
+
+def sidak_z_for(count):
+ # per-test z holding the family-wise confidence of verdict_z across count candidates
+ # (sidak): the selected best of several noisy rows sits low by selection, so testing
+ # it at the single-test z would overstate confidence
+ if count <= 1:
+ return verdict_z
+ single = statistics.NormalDist().cdf(verdict_z)
+ return statistics.NormalDist().inv_cdf(single ** (1.0 / count))
# the synthetic-tensor errors are seed-fixed math and land within a few percent across
# nvidia, amd and intel gpus; a value outside its band means the measurement itself is
@@ -2527,6 +2536,58 @@ def best_config(results):
return min(near_fastest, key=lambda item: item[2])[0]
+def config_settings(kwargs):
+ # ui settings tuple a bench config corresponds to; None for the external baselines
+ # and the unquantized row, which are not reachable states of the attention dropdowns
+ if not kwargs or not kwargs.get("do_quantize", True):
+ return None
+ matmul = kwargs.get("matmul_dtype", "auto")
+ pv = kwargs.get("pv_matmul_dtype", "auto")
+ return {
+ "matmul": "enabled" if matmul == "auto" else matmul,
+ "pv": "disabled" if pv == "auto" else pv,
+ "smooth": bool(kwargs.get("smooth_k", False)),
+ "hadamard": bool(kwargs.get("use_hadamard", False)),
+ }
+
+
+def select_attention_config(results):
+ # the attention settings jointly pick one kernel config, so they are judged as one:
+ # rank the measured quantized rows by the same rule as the per-shape star instead of
+ # testing each setting alone against int8, which can assemble a settings tuple no
+ # row ever measured. rows breaching the incremental error cap vs int8 are held out
+ # of the ranking but kept for citation
+ int8_err_ref = measured(results, "int8")[1]
+ pool, capped = [], []
+ for config_id, label, kwargs in bench_configs:
+ settings = config_settings(kwargs)
+ if settings is None:
+ continue
+ ms, err = measured(results, config_id)
+ if ms is None:
+ continue
+ entry = dict(config_id=config_id, label=label.removeprefix("sdnq "), ms=ms, err=err, settings=settings)
+ if err and int8_err_ref and err > int8_err_ref * recommend_error_cap:
+ capped.append(entry)
+ else:
+ pool.append(entry)
+ if not pool:
+ return None, pool, capped
+ fastest = min(entry["ms"] for entry in pool)
+ window = [entry for entry in pool if entry["ms"] <= fastest * 1.05]
+ chosen = min(window, key=lambda entry: entry["err"] if entry["err"] is not None else float("inf"))
+ return chosen, pool, capped
+
+
+def best_in_subset(entries):
+ # the star rule inside one setting's subset, for citing its strongest variant
+ if not entries:
+ return None
+ fastest = min(entry["ms"] for entry in entries)
+ window = [entry for entry in entries if entry["ms"] <= fastest * 1.05]
+ return min(window, key=lambda entry: entry["err"] if entry["err"] is not None else float("inf"))
+
+
def build_recommendations(all_results, fp8_result, prep_status, block_results=None, block_variants=None):
# prefer an image dit shape with the full config set as reference, then the video shapes
reference = None
@@ -2552,34 +2613,35 @@ def build_recommendations(all_results, fp8_result, prep_status, block_results=No
return str(getattr(shared.opts, key))
rows = []
- # verdicts are computed before any row is built: the smooth k and hadamard buybacks
- # feed a composition check on the qk verdict, since toggles that pass individually
- # can still lose to unquantized as a stack. compare against the unquantized sdnq row:
- # quantization has to pay for its own prep and clear the shared speed margin, at this
- # run's own measured noise level; too close to call keeps the current setting
- qk_sigma = pair_sigma(results.get("int8") or {}, "ms", results.get("noquant") or {}, "ms")
- qk_test = speed_verdict(int8_ms, noquant_ms, sigma=qk_sigma)
- use_quantized = qk_test == "faster"
- qk_inconclusive = qk_test == "inconclusive"
- if int8_ms and noquant_ms:
- if use_quantized:
- quant_reason = f"int8 qk measured x{noquant_ms / int8_ms:.2f} vs unquantized sdnq attention"
- if int8_err and noquant_err:
- quant_reason += f", error {int8_err:.5f} vs {noquant_err:.5f}; smooth k and hadamard below buy error back"
- elif qk_inconclusive:
- quant_reason = f"int8 qk measured x{noquant_ms / int8_ms:.2f} vs unquantized sdnq attention, too close to the x{1 / recommend_speed_margin:.2f} margin to call at this run's noise (ratio sigma {qk_sigma:.1%}); keeping the current setting"
- elif int8_ms < noquant_ms:
- quant_reason = f"int8 qk measured only x{noquant_ms / int8_ms:.2f} vs unquantized sdnq attention, under the x{1 / recommend_speed_margin:.2f} margin the verdict requires"
- if int8_err and noquant_err:
- quant_reason += f"; unquantized keeps error {noquant_err:.5f} vs {int8_err:.5f}"
+ # the settings below jointly pick one kernel config, so the verdict is joint: select
+ # the best measured quantized row (star rule with the error cap), then gate that one
+ # config against the unquantized sdnq row: quantization has to pay for its own prep
+ # and clear the shared speed margin, at this run's own measured noise level and at a
+ # z adjusted for having selected the best of the pool; too close to call keeps the
+ # current settings
+ chosen, pool, capped = select_attention_config(results)
+ gate = None
+ quant_reason = "no quantized attention config ran at this shape"
+ if chosen is not None and noquant_ms:
+ gate_sigma = pair_sigma(results.get(chosen["config_id"]) or {}, "ms", results.get("noquant") or {}, "ms")
+ gate = speed_verdict(chosen["ms"], noquant_ms, sigma=gate_sigma, z=sidak_z_for(len(pool)))
+ if gate == "faster":
+ quant_reason = f"{chosen['label']} measured x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention"
+ if chosen["err"] and noquant_err:
+ quant_reason += f", error {chosen['err']:.5f} vs {noquant_err:.5f}"
+ elif gate == "inconclusive":
+ quant_reason = f"best quantized config ({chosen['label']}) measured x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention, too close to the x{1 / recommend_speed_margin:.2f} margin to call at this run's noise (ratio sigma {gate_sigma:.1%}); keeping the current setting"
+ elif chosen["ms"] < noquant_ms:
+ quant_reason = f"best quantized config ({chosen['label']}) measured only x{noquant_ms / chosen['ms']:.2f} vs unquantized sdnq attention, under the x{1 / recommend_speed_margin:.2f} margin the verdict requires"
+ if chosen["err"] and noquant_err:
+ quant_reason += f"; unquantized keeps error {noquant_err:.5f} vs {chosen['err']:.5f}"
else:
- quant_reason = f"unquantized sdnq measured x{int8_ms / noquant_ms:.2f} vs int8 qk with lower error; quantization prep outweighs the kernel gain on this gpu"
- elif int8_ms and base_ms:
- use_quantized = base_ms / int8_ms >= 1.10
- quant_reason = f"int8 qk measured x{base_ms / int8_ms:.2f} vs torch sdpa; unquantized sdnq row unavailable"
- else:
- use_quantized = False
- quant_reason = "int8 qk failed to run"
+ quant_reason = f"unquantized sdnq measured x{chosen['ms'] / noquant_ms:.2f} vs the best quantized config ({chosen['label']}) with lower error; quantization prep outweighs the kernel gain on this gpu"
+ elif chosen is not None and base_ms:
+ gate = "faster" if base_ms / chosen["ms"] >= 1.10 else "not_faster"
+ quant_reason = f"{chosen['label']} measured x{base_ms / chosen['ms']:.2f} vs torch sdpa; unquantized sdnq row unavailable"
+ use_quantized = gate == "faster"
+ qk_inconclusive = gate == "inconclusive"
# buyback costs are judged at the block geometry matching the reference family when
# one was measured, so a krea2 verdict uses the krea2-width block
@@ -2623,127 +2685,105 @@ def build_recommendations(all_results, fp8_result, prep_status, block_results=No
hadamard_rec = hadamard_gain >= 1.3 and hadamard_cost <= 0.15
hadamard_reason = f"int8 error x{hadamard_gain:.1f} lower for {hadamard_cost_note}; hangs torch compile on non pow2 head dims (SD 1.5)"
- # composition check: the toggles recommended above must still beat unquantized as a
- # stack; when the exact stack was not measured, predict it additively from the single
- # toggle deltas (measured cross-vendor: additive to ~2% median with a +2% skew, so the
- # prediction carries that correction plus a model sigma alongside the measurement noise)
- additive_model_skew = 1.02
- additive_model_sigma = 0.028
+ # rows decompose the one selected config; each reason cites the measured row that
+ # isolates its own setting where one was benchmarked
+ by_settings = {}
+ for entry in pool + capped:
+ s = entry["settings"]
+ by_settings[(s["matmul"], s["pv"], s["smooth"], s["hadamard"])] = entry
- def stack_estimate(qk_ms, toggles):
- deltas, measured_row = 0.0, results.get({(True, True): "smooth_hadamard", (True, False): "smooth", (False, True): "hadamard", (False, False): "int8"}[toggles]) or {}
- if measured_row.get("ms"):
- return measured_row["ms"], row_sigma(measured_row, "ms"), False
- for on, toggle_id in ((toggles[0], "smooth"), (toggles[1], "hadamard")):
- toggle_ms, _err = measured(results, toggle_id)
- if on and toggle_ms and int8_ms:
- deltas += toggle_ms - int8_ms
- sigma = row_sigma(results.get("int8") or {}, "ms")
- sigma = math.sqrt((sigma or 0.0) ** 2 + additive_model_sigma ** 2)
- return (qk_ms + deltas) * additive_model_skew, sigma, True
+ def sibling(entry, **overrides):
+ s = dict(entry["settings"], **overrides)
+ other = by_settings.get((s["matmul"], s["pv"], s["smooth"], s["hadamard"]))
+ return None if other is entry else other
- composition_flip = False
- if use_quantized and noquant_ms and int8_ms:
- toggles = (bool(smooth_rec), bool(hadamard_rec))
- combo_label = {(True, True): "int8 qk + smooth + hadamard", (True, False): "int8 qk + smooth k", (False, True): "int8 qk + hadamard", (False, False): "int8 qk"}[toggles]
- combo_ms, combo_sigma, estimated = stack_estimate(int8_ms, toggles)
- combo_sigma = math.sqrt((combo_sigma or 0.0) ** 2 + (row_sigma(results.get("noquant") or {}, "ms") or 0.0) ** 2) or None
- stack_test = speed_verdict(combo_ms, noquant_ms, sigma=combo_sigma)
- if combo_ms and stack_test != "faster":
- use_quantized = False
- composition_flip = True
- source = "estimated additively at" if estimated else "measured"
- if stack_test == "inconclusive":
- quant_reason = f"the recommended stack ({combo_label}) {source} x{noquant_ms / combo_ms:.2f} vs unquantized sdnq attention, too close to the margin to call; disabled until it clearly wins"
- else:
- quant_reason = f"the recommended stack ({combo_label}) {source} x{noquant_ms / combo_ms:.2f} vs unquantized sdnq attention; the error buybacks eat the qk gain on this gpu"
+ def block_buyback_survives(block_config_id):
+ # None when the block section did not measure the pair; False when the toggle
+ # left block output error unchanged, so its kernel-scope buyback is cosmetic here
+ variant_err = (block_rows.get(block_config_id) or {}).get("err")
+ base_block_err = (block_rows.get("int8-mm-atten") or {}).get("err")
+ if not (variant_err and base_block_err):
+ return None
+ return variant_err < base_block_err * 0.98
- # the verdict must not hinge on int8 alone: bare float8 qk can clear the margin on
- # gpus where int8 falls short, so it gets its own shot before disabling
fp8qk_ms, fp8qk_err = measured(results, "fp8qk")
- fp8_rescue = False
- if not use_quantized and not composition_flip and noquant_ms and fp8qk_ms:
- fp8_sigma = pair_sigma(results.get("fp8qk") or {}, "ms", results.get("noquant") or {}, "ms")
- if speed_verdict(fp8qk_ms, noquant_ms, sigma=fp8_sigma) == "faster" and not (fp8qk_err and int8_err and fp8qk_err > int8_err * recommend_error_cap):
- use_quantized = True
- fp8_rescue = True
- qk_inconclusive = False
- quant_reason = f"float8 qk measured x{noquant_ms / fp8qk_ms:.2f} vs unquantized sdnq attention where the int8 path fell short"
- if fp8qk_err:
- quant_reason += f", error {fp8qk_err:.5f}"
- quant_reason += "; smooth k and hadamard buybacks were only measured on int8"
-
- qk_choice = "enabled"
- qk_reason = quant_reason + "; enabled resolves to int8, uint8 remaps to int8"
- if fp8_rescue:
- qk_choice = "float8_e4m3fn"
+ if use_quantized:
+ qk_choice = chosen["settings"]["matmul"]
qk_reason = quant_reason
- elif fp8qk_ms and int8_ms:
- qk_compare = f"float8 qk measured x{int8_ms / fp8qk_ms:.2f} vs int8"
- if fp8qk_err and int8_err:
- qk_compare += f", error {fp8qk_err:.5f} vs {int8_err:.5f}"
- if fp8qk_ms < int8_ms * 0.95 and not (fp8qk_err and int8_err and fp8qk_err > int8_err * recommend_error_cap):
- qk_choice = "float8_e4m3fn"
- qk_reason = quant_reason + "; " + qk_compare
- else:
- qk_reason += f"; {qk_compare}"
- elif fp8_result["qk"][0]:
- qk_reason += "; float8 compiles here but was not benchmarked at this shape"
- if qk_inconclusive:
+ faster_capped = [entry for entry in capped if entry["ms"] < chosen["ms"]]
+ if faster_capped and int8_err:
+ fastest_capped = min(faster_capped, key=lambda entry: entry["ms"])
+ qk_reason += f"; {fastest_capped['label']} is x{chosen['ms'] / fastest_capped['ms']:.2f} faster but multiplies error x{fastest_capped['err'] / int8_err:.1f} over int8"
+ if qk_choice == "enabled":
+ qk_reason += "; enabled resolves to int8, uint8 remaps to int8"
+ if fp8qk_ms and int8_ms and qk_choice != "float8_e4m3fn":
+ qk_reason += f"; float8 qk measured x{int8_ms / fp8qk_ms:.2f} vs int8"
+ if fp8qk_err and int8_err:
+ qk_reason += f", error {fp8qk_err:.5f} vs {int8_err:.5f}"
+ elif fp8_result["qk"][0] and not fp8qk_ms:
+ qk_reason += "; float8 compiles here but was not benchmarked at this shape"
+ elif qk_inconclusive:
qk_choice = current("sdnq_attention_matmul_type")
qk_reason = quant_reason
- elif not use_quantized:
+ else:
qk_choice = "disabled"
qk_reason = quant_reason
rows.append(("MatMul type", current("sdnq_attention_matmul_type"), qk_choice, qk_reason))
- # pv rows were measured on top of int8 qk, so a pv verdict only holds when qk
- # quantization itself is recommended; note enabled means int8 pv on this dropdown.
- # each candidate is tested against the margin independently at a sidak-adjusted z:
- # picking the fastest first and testing it afterwards would bias toward enabling,
- # since the minimum of several noisy rows sits low by selection
- pv_candidates = []
- for pv_dtype, pv_name, pv_id in (("float8_e4m3fn", "fp8", "fp8pv"), ("int8", "int8", "int8pv"), ("float16", "fp16", "fp16pv")):
- pv_ms, pv_err = measured(results, pv_id)
- if pv_ms:
- pv_candidates.append((pv_dtype, pv_name, pv_id, pv_ms, pv_err))
- if use_quantized and pv_candidates and int8_ms:
- z_sel = sidak_z.get(len(pv_candidates), sidak_z[4])
- pv_winners, pv_inconclusive = [], False
- for pv_dtype, pv_name, pv_id, pv_ms, pv_err in pv_candidates:
- pv_test = speed_verdict(pv_ms, int8_ms, sigma=pair_sigma(results.get(pv_id) or {}, "ms", results.get("int8") or {}, "ms"), z=z_sel)
- if pv_test == "faster":
- pv_winners.append((pv_dtype, pv_name, pv_ms, pv_err))
- elif pv_test == "inconclusive":
- pv_inconclusive = True
- if pv_winners:
- fastest_pv = min(ms for _d, _n, ms, _e in pv_winners)
- near_fastest = [c for c in pv_winners if c[2] <= fastest_pv * 1.05]
- pv_dtype, pv_name, pv_ms, pv_err = min(near_fastest, key=lambda c: c[3] if c[3] is not None else float("inf"))
- if pv_err and int8_err and pv_err > int8_err * recommend_error_cap:
- rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), "disabled", f"{pv_name} pv is x{int8_ms / pv_ms:.2f} faster but multiplies error x{pv_err / int8_err:.1f}; disabled keeps pv unquantized"))
- else:
- pv_reason = f"{pv_name} pv measured x{int8_ms / pv_ms:.2f} over int8 qk alone"
- if pv_err and int8_err:
- pv_reason += f", error {pv_err:.5f} vs {int8_err:.5f}"
- rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), pv_dtype, pv_reason))
- elif pv_inconclusive:
- rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), current("sdnq_attention_pv_matmul_type"), "the best pv candidate sits within this run's noise of the margin; keeping the current setting"))
+ pv_names = {"int8": "int8", "float16": "fp16", "float8_e4m3fn": "fp8"}
+ if use_quantized:
+ pv_choice = chosen["settings"]["pv"]
+ if pv_choice != "disabled":
+ without = sibling(chosen, pv="disabled")
+ pv_reason = f"{pv_names[pv_choice]} pv rides the selected config"
+ if without:
+ pv_reason = f"{pv_names[pv_choice]} pv measured x{without['ms'] / chosen['ms']:.2f} over the same stack without pv"
+ if chosen["err"] and without["err"]:
+ pv_reason += f", error {chosen['err']:.5f} vs {without['err']:.5f}"
else:
- rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), "disabled", f"disabled keeps pv unquantized; {' and '.join(name for _d, name, _i, _m, _e in pv_candidates)} pv measured no gain here"))
+ best_pv = best_in_subset([entry for entry in pool + capped if entry["settings"]["pv"] != "disabled"])
+ pv_reason = "disabled keeps pv unquantized"
+ if best_pv is not None:
+ pv_label = f"{pv_names[best_pv['settings']['pv']]} pv ({best_pv['label']})"
+ if best_pv["err"] and int8_err and best_pv["err"] > int8_err * recommend_error_cap:
+ pv_reason = f"{pv_label} is x{chosen['ms'] / best_pv['ms']:.2f} the speed of the selected config but multiplies error x{best_pv['err'] / int8_err:.1f} over int8; disabled keeps pv unquantized"
+ else:
+ pv_reason += f"; {best_pv['label']} measured x{chosen['ms'] / best_pv['ms']:.2f} the speed of the selected config"
+ if best_pv["err"] and chosen["err"]:
+ pv_reason += f", error {best_pv['err']:.5f} vs {chosen['err']:.5f}"
+ rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), pv_choice, pv_reason))
+ elif qk_inconclusive:
+ rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), current("sdnq_attention_pv_matmul_type"), "the qk verdict above is inconclusive; pv follows it"))
else:
- if qk_inconclusive and pv_candidates:
- pv_note = "the qk verdict above is inconclusive; pv follows it"
- elif not use_quantized and pv_candidates:
- pv_note = "qk quantization is not recommended above; pv on an unquantized qk path was not measured"
- else:
- pv_note = "disabled keeps pv unquantized"
- rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), "disabled" if not qk_inconclusive else current("sdnq_attention_pv_matmul_type"), pv_note))
+ rows.append(("PV MatMul type", current("sdnq_attention_pv_matmul_type"), "disabled", "qk quantization is not recommended above; pv on an unquantized qk path was not measured"))
- if smooth_rec is not None:
- rows.append(("Use Smooth K", current("sdnq_attention_smooth_k"), str(smooth_rec), smooth_reason))
- if hadamard_rec is not None:
- rows.append(("Use Hadamard", current("sdnq_attention_use_hadamard"), str(hadamard_rec), hadamard_reason))
+ def toggle_row(setting_key, ui_name, on, buyback_reason, subset, block_config_id, caveat=None):
+ # a toggle the selected config carries keeps the buyback evidence as its reason;
+ # one it omits cites the strongest variant that carried it, plus the block-scope
+ # cross-check when that shows the buyback never reached block output
+ if on:
+ reason = buyback_reason or f"part of the selected config ({chosen['label']})"
+ if caveat and caveat not in reason:
+ reason += f"; {caveat}"
+ else:
+ reason = "the selected config omits it"
+ best_variant = best_in_subset(subset)
+ if best_variant is not None:
+ reason += f"; {best_variant['label']} measured x{chosen['ms'] / best_variant['ms']:.2f} the speed of the selected config"
+ if best_variant["err"] and chosen["err"]:
+ reason += f", error {best_variant['err']:.5f} vs {chosen['err']:.5f}"
+ if block_buyback_survives(block_config_id) is False:
+ reason += "; the block cross-check found no output-error change from it at this shape"
+ rows.append((ui_name, current(setting_key), str(bool(on)), reason))
+
+ if use_quantized:
+ toggle_row("sdnq_attention_smooth_k", "Use Smooth K", chosen["settings"]["smooth"], smooth_reason, [entry for entry in pool + capped if entry["settings"]["smooth"]], "int8-mm-smooth")
+ toggle_row("sdnq_attention_use_hadamard", "Use Hadamard", chosen["settings"]["hadamard"], hadamard_reason, [entry for entry in pool + capped if entry["settings"]["hadamard"]], "int8-mm-hadamard", caveat="hangs torch compile on non pow2 head dims (SD 1.5)")
+ else:
+ if smooth_rec is not None:
+ rows.append(("Use Smooth K", current("sdnq_attention_smooth_k"), str(smooth_rec), smooth_reason))
+ if hadamard_rec is not None:
+ rows.append(("Use Hadamard", current("sdnq_attention_use_hadamard"), str(hadamard_rec), hadamard_reason))
rows.append(("Hadamard Group Size", current("sdnq_attention_hadamard_group_size"), "256", "values above head dim are clamped; non pow2 values floor to the nearest power of 2"))
@@ -2761,27 +2801,29 @@ def build_recommendations(all_results, fp8_result, prep_status, block_results=No
emit(f"[yellow]{differing} settings differ from the recommended values, change them in Compute Settings -> SDNQ Attention[/yellow]")
else:
emit("[green]current settings already match the recommendations[/green]")
- if not use_quantized:
+ if not use_quantized and not qk_inconclusive:
emit("[dim]qk quantization is not worth it here, so MatMul type recommends disabled; the smooth k and hadamard rows show what to pick if it is enabled anyway[/dim]")
notes = []
- # per-shape qk verdicts: the table above judges one reference shape, but the settings
- # are global and workloads differ; a split gpu (video wins, image loses) shows here
- # rather than being averaged away. cross and te shapes are prep-dominated, skip them
+ # per-shape joint verdicts: the table above judges one reference shape, but the
+ # settings are global and workloads differ; a split gpu (video wins, image loses)
+ # shows here rather than being averaged away, and each shape names the config its
+ # star selected. cross and te shapes are prep-dominated, skip them
shape_verdicts = []
for shape_label in recommendation_presets:
shape_results = all_results.get(shape_label)
if not shape_results:
continue
shape_noquant_ms, _err = measured(shape_results, "noquant")
- shape_int8_ms, _err = measured(shape_results, "int8")
- if shape_noquant_ms and shape_int8_ms:
- shape_test = speed_verdict(shape_int8_ms, shape_noquant_ms, sigma=pair_sigma(shape_results.get("int8") or {}, "ms", shape_results.get("noquant") or {}, "ms"))
+ shape_chosen, shape_pool, _shape_capped = select_attention_config(shape_results)
+ if shape_noquant_ms and shape_chosen is not None:
+ shape_sigma = pair_sigma(shape_results.get(shape_chosen["config_id"]) or {}, "ms", shape_results.get("noquant") or {}, "ms")
+ shape_test = speed_verdict(shape_chosen["ms"], shape_noquant_ms, sigma=shape_sigma, z=sidak_z_for(len(shape_pool)))
word = {"faster": "enabled", "not_faster": "disabled", "inconclusive": "inconclusive"}[shape_test]
- shape_verdicts.append((word, f"{shape_label} {word} (x{shape_noquant_ms / shape_int8_ms:.2f})"))
+ shape_verdicts.append((word, f"{shape_label} {word} ({shape_chosen['label']} x{shape_noquant_ms / shape_chosen['ms']:.2f})"))
if len(shape_verdicts) > 1:
split = len({word for word, _text in shape_verdicts}) > 1
- line = f"qk verdict by shape: {', '.join(text for _word, text in shape_verdicts)}"
+ line = f"verdict by shape: {', '.join(text for _word, text in shape_verdicts)}"
if split:
notes.append(f"[yellow]{line}; settings are global, pick for the shapes you generate at[/yellow]")
else:
@@ -2959,7 +3001,7 @@ def build_dequant_recommendations(dequant_results, weight_dequant_result, varian
fwd_err = mm_entry.get("fwd_err")
err_ok = not (fwd_err and best_err) or best_err <= fwd_err * recommend_error_cap
best_entry = mm_entry if best_sel == "enabled" else float_mm_entry(mm_id, best_sel)
- mm_test = speed_verdict(best_ms, mm_entry["fwd_ms"], sigma=pair_sigma(best_entry, "mm_ms", mm_entry, "fwd_ms"), z=sidak_z.get(len(mm_candidates), sidak_z[4]))
+ mm_test = speed_verdict(best_ms, mm_entry["fwd_ms"], sigma=pair_sigma(best_entry, "mm_ms", mm_entry, "fwd_ms"), z=sidak_z_for(len(mm_candidates)))
recommend_mm = mm_test == "faster" and err_ok
mm_reason = f"{mm_id} quantized matmul ({best_resolved}) measured {ratio_text(mm_entry['fwd_ms'], best_ms)} vs the dequant path"
if fwd_err and best_err:
@@ -3084,7 +3126,7 @@ def build_dequant_recommendations(dequant_results, weight_dequant_result, varian
rows.append(("Quantize convolutional layers", current("sdnq_quantize_conv_layers"), str(conv_free), conv_reason))
if conv_int8.get("fwd_ms") and conv_mm_rows:
best_mm_id, best_mm = min(conv_mm_rows, key=lambda item: item[1]["fwd_ms"])
- conv_mm_test = speed_verdict(best_mm["fwd_ms"], conv_int8["fwd_ms"], sigma=pair_sigma(best_mm, "fwd_ms", conv_int8, "fwd_ms"), z=sidak_z.get(len(conv_mm_rows), sidak_z[4]))
+ conv_mm_test = speed_verdict(best_mm["fwd_ms"], conv_int8["fwd_ms"], sigma=pair_sigma(best_mm, "fwd_ms", conv_int8, "fwd_ms"), z=sidak_z_for(len(conv_mm_rows)))
mm_err_ok = not (best_mm.get("out_err") and conv_int8.get("out_err")) or best_mm["out_err"] <= conv_int8["out_err"] * recommend_error_cap
conv_mm_reason = f"{best_mm_id} measured {best_mm['fwd_ms']:.3f} vs {conv_int8['fwd_ms']:.3f} ms for the int8 conv dequant path at {conv_shapes[0][0]}, output error {best_mm.get('out_err', 0):.5f} vs {conv_int8.get('out_err', 0):.5f}"
other_base, _other_int8, other_mm_rows = conv_per_shape.get(conv_shapes[1][0], ({}, {}, [])) if len(conv_shapes) > 1 else ({}, {}, [])
diff --git a/data/previews.json b/data/previews.json
index 30ca5de77..41a9282d3 100644
--- a/data/previews.json
+++ b/data/previews.json
@@ -30,7 +30,7 @@
"THUDM--CogVideoX-5b": "THUDM--CogView3-Plus-3B.jpg",
"vladmandic--Anima-1.0-Base-Merge-sdnq-hadamard-uint4": "vladmandic--Anima-1.0-Base.jpg",
"vladmandic--Anima-1.0-Base-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Base.jpg",
- "vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Turbo.jpg",
+ "vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4": "vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4.jpg",
"vladmandic--Flux.2-Klein-9B-KV-sdnq-hadamard-uint4": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
"vladmandic--Flux.2-Klein-9B-KV-Merge-sdnq-hadamard-uint4": "black-forest-labs--FLUX.2-klein-9b-kv.jpg",
"vladmandic--Krea-2-Base-sdnq-hadamard-uint4": "CalamitousFelicitousness--Krea-2-Base-Diffusers.jpg",
diff --git a/data/reference-base.json b/data/reference-base.json
index 70c3d8b58..96b18d90d 100644
--- a/data/reference-base.json
+++ b/data/reference-base.json
@@ -189,7 +189,7 @@
"path": "CalamitousFelicitousness/Krea-2-Base-Diffusers",
"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, width: 1024, height: 1024",
+ "extras": "sampler: Default, cfg_scale: 4.5, steps: 52",
"size": 33.5,
"date": "2026 June"
},
@@ -918,5 +918,45 @@
"size": 47.98,
"extras": "sampler: Default",
"date": "2026 July"
+ },
+ "Microsoft Mage-Flow": {
+ "path": "vladmandic/Mage-Flow-4B",
+ "preview": "vladmandic--Mage-Flow-4B.jpg",
+ "desc": "Mage-Flow is a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing.",
+ "extras": "sampler: Default",
+ "size": 16.19,
+ "date": "2026 July"
+ },
+ "SeFi-Image 1B Base": {
+ "path": "SeFi-Image/SeFi-Image-1B-Base-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-1B-Base-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 6.32,
+ "date": "2026 July"
+ },
+ "SeFi-Image 2B Base": {
+ "path": "SeFi-Image/SeFi-Image-2B-Base-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-2B-Base-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 8.18,
+ "date": "2026 July"
+ },
+ "SeFi-Image 5B Base": {
+ "path": "SeFi-Image/SeFi-Image-5B-Base-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-5B-Base-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 17.69,
+ "date": "2026 July"
+ },
+ "SeFi-Image 5B RL": {
+ "path": "SeFi-Image/SeFi-Image-5B-RL-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-5B-RL-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 17.69,
+ "date": "2026 July"
}
}
diff --git a/data/reference-distilled.json b/data/reference-distilled.json
index f1391ec63..96f70d453 100644
--- a/data/reference-distilled.json
+++ b/data/reference-distilled.json
@@ -26,7 +26,7 @@
"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, width: 1024, height: 1024",
+ "extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 33.5,
"date": "2026 June"
},
@@ -129,7 +129,7 @@
"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": "width: 2048, height: 1024, sampler: DEIS, steps: 40, cfg_scale: 6.0",
+ "extras": "sampler: DEIS, steps: 40, cfg_scale: 6.0",
"experimental": true
},
"NVLabs Sana 1.5 1.6B 1k Sprint": {
@@ -225,5 +225,37 @@
"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",
+ "desc": "Mage-Flow is a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing.",
+ "extras": "sampler: Default",
+ "size": 16.19,
+ "date": "2026 July"
+ },
+ "SeFi-Image 1B Turbo": {
+ "path": "SeFi-Image/SeFi-Image-1B-turbo-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-1B-turbo-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 6.32,
+ "date": "2026 July"
+ },
+ "SeFi-Image 2B Turbo": {
+ "path": "SeFi-Image/SeFi-Image-2B-turbo-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-2B-turbo-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 8.18,
+ "date": "2026 July"
+ },
+ "SeFi-Image 5B Turbo": {
+ "path": "SeFi-Image/SeFi-Image-5B-turbo-diffusers",
+ "preview": "SeFi-Image--SeFi-Image-5B-turbo-diffusers.jpg",
+ "desc": "SeFi-Image is a text-to-image foundation model family built with Semantic-First Diffusion. It separates generation into semantic and texture latent streams, denoising semantic structure slightly ahead of texture details.",
+ "extras": "sampler: Default",
+ "size": 17.69,
+ "date": "2026 July"
}
}
diff --git a/data/reference-nunchaku.json b/data/reference-nunchaku.json
index 002a6bb46..cffe2ecd5 100644
--- a/data/reference-nunchaku.json
+++ b/data/reference-nunchaku.json
@@ -212,5 +212,67 @@
"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",
+ "extras": "sampler: Default, cfg_scale: 1.0, steps: 9",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 6.17,
+ "date": "2026 July"
+ },
+ "Baidu ERNIE-Image Nunchaku-Lite": {
+ "path": "lite-infer/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder",
+ "preview": "baidu--ERNIE-Image.jpg",
+ "extras": "sampler: Default, cfg_scale: 4.0, steps: 50",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 6.88,
+ "date": "2026 July"
+ },
+ "BFL FLUX.1 Dev Nunchaku-Lite": {
+ "path": "lite-infer/flux.1-dev-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "black-forest-labs--FLUX.1-dev.jpg",
+ "extras": "sampler: Default, cfg_scale: 3.5",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 10.97,
+ "date": "2026 July"
+ },
+ "BFL FLUX.1 Schnell Nunchaku-Lite": {
+ "path": "lite-infer/flux.1-schnell-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "black-forest-labs--FLUX.1-schnell.jpg",
+ "extras": "sampler: Default, cfg_scale: 3.5",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 10.95,
+ "date": "2026 July"
+ },
+ "BFL FLUX.1 Kontext Nunchaku-Lite": {
+ "path": "lite-infer/flux.1-kontext-dev-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "black-forest-labs--FLUX.1-Kontext-dev.jpg",
+ "extras": "sampler: Default, cfg_scale: 3.5",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 10.97,
+ "date": "2026 July"
+ },
+ "BFL FLUX.1 Krea Nunchaku-Lite": {
+ "path": "lite-infer/flux.1-krea-dev-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "black-forest-labs--FLUX.1-Krea-dev.jpg",
+ "extras": "sampler: Default, cfg_scale: 4.5",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 10.97,
+ "date": "2026 July"
+ },
+ "Qwen-Image Nunchaku-Lite": {
+ "path": "lite-infer/Qwen-Image-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "Qwen--Qwen-Image.jpg",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 16.81,
+ "date": "2026 July"
+ },
+ "Qwen-Image-Edit-2509 Nunchaku-Lite": {
+ "path": "lite-infer/qwen-image-edit-2509-nunchaku-lite-int4_r32-bnb4-text-encoder",
+ "preview": "Qwen--Qwen-Image-Edit-2509.jpg",
+ "desc": "Nunchaku-Lite quantization using precompiled Kernels",
+ "size": 16.81,
+ "date": "2026 July"
}
}
diff --git a/data/reference-quantized.json b/data/reference-quantized.json
index c544542a6..132fdfe6e 100644
--- a/data/reference-quantized.json
+++ b/data/reference-quantized.json
@@ -197,7 +197,7 @@
},
"Anima 1.0 Turbo sdnq-svd-dynamic-uint4": {
"path": "vladmandic/Anima-1.0-Turbo-sdnq-svd-dynamic-uint4",
- "preview": "vladmandic--Anima-1.0-Turbo.jpg",
+ "preview": "vladmandic--Anima-1.0-Turbo-sdnq-svd-dynamic-uint4.jpg",
"desc": "Anima 1.0 Turbo with extended 1024-resolution training and expanded dataset coverage for less common artists. A 2B parameter anime-focused text-to-image model based on modified Cosmos-Predict-2B with Qwen3-0.6B text encoder, created by CircleStone Labs and Comfy Org.",
"date": "2026 May",
"size": 2.03
@@ -230,7 +230,7 @@
"path": "vladmandic/Krea-2-Turbo-sdnq-hadamard-uint4",
"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, width: 1024, height: 1024",
+ "extras": "sampler: Default, cfg_scale: 1.0, steps: 8",
"size": 10.54,
"date": "2026 July"
},
@@ -238,7 +238,7 @@
"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, width: 1024, height: 1024",
+ "extras": "sampler: Default, cfg_scale: 4.5, steps: 52",
"size": 10.3,
"date": "2026 June"
}
diff --git a/extensions-builtin/sdnext-kanvas b/extensions-builtin/sdnext-kanvas
index dc47fa212..b98510429 160000
--- a/extensions-builtin/sdnext-kanvas
+++ b/extensions-builtin/sdnext-kanvas
@@ -1 +1 @@
-Subproject commit dc47fa2129673caa69ffadeeffba49700e1a7d4a
+Subproject commit b985104298cb37717c5ff34c59e0b186cb40757e
diff --git a/extensions-builtin/sdnext-modernui b/extensions-builtin/sdnext-modernui
index 45f0e695e..20f032f60 160000
--- a/extensions-builtin/sdnext-modernui
+++ b/extensions-builtin/sdnext-modernui
@@ -1 +1 @@
-Subproject commit 45f0e695eb05527a13d78e41a8e327cd8ea176a5
+Subproject commit 20f032f601bf37097f093e2b184070859931d810
diff --git a/installer.py b/installer.py
index 8507ac57e..cef9feb63 100644
--- a/installer.py
+++ b/installer.py
@@ -127,8 +127,13 @@ def env_flag(name: str, default: bool = False) -> bool:
def print_profile(profiler: cProfile.Profile, msg: str):
profiler.disable()
- from modules.errors import profile
- profile(profiler, msg)
+ from modules.errors import profile_print
+ profile_print(msg, local_profiler=profiler)
+
+
+def profile(*_args, **_kwargs):
+ # legacy to avoid import errors
+ pass
def package_version(package):
@@ -549,7 +554,7 @@ def check_diffusers():
t_start = time.time()
if args.skip_all:
return
- target_commit = "01969142b55379991fee07608c9e7e8f80afced0" # diffusers commit hash == 0.39.0.dev0 == 06-29-2026
+ target_commit = "6f2010e8bbe61fd2a81a659b858e298edcba8fab" # diffusers commit hash == 0.40.0.dev0 == 08-04-2026
# if args.use_rocm or args.use_zluda or args.use_directml:
# sha = '043ab2520f6a19fce78e6e060a68dbc947edb9f9' # lock diffusers versions for now
pkg = package_spec('diffusers')
@@ -560,7 +565,7 @@ def check_diffusers():
if minor == -1:
log.info(f'Install: package="diffusers" commit={target_commit}')
else:
- log.info(f'Update: package="diffusers" current={pkg.version} hash={current} target={target_commit}')
+ log.info(f'Update: package="diffusers" current={pkg.version} commit={current} target={target_commit}')
pip('uninstall --yes diffusers', ignore=True, quiet=True, uv=False)
if args.skip_git:
log.warning('Git: marked as not available but required for diffusers installation')
@@ -578,7 +583,7 @@ def check_transformers():
pkg_transformers = package_spec('transformers')
pkg_tokenizers = package_spec('tokenizers')
# target_commit = '753d61104116eefc8ffc977327b441ee0c8d599f' # transformers commit hash == 4.57.6
- # target_commit = "380e3cc5d59912a48508cb6d4959a31cd460e12e" # transformers commit hash == 5.5.0.dev-0409
+ # target_commit = "cf8572d34e39818e42dbf220701fbd3eb5b5a82a" # transformers commit hash == 5.14.0.dev0 == 08-04-2026
target_commit = "b70d02fc724d04c916832ca4ead03ff05e8fb1ee" # transformers commit hash == 5.13.0.dev0 == 07-03-2026
if args.use_directml:
target_transformers = '4.52.4'
@@ -604,7 +609,7 @@ def check_transformers():
if pkg_transformers is None:
log.info(f'Install: package="transformers" commit={target_commit}')
else:
- log.info(f'Update: package="transformers" current={pkg_transformers.version} hash={current} target={target_commit}')
+ log.info(f'Update: package="transformers" current={pkg_transformers.version} commit={current} target={target_commit}')
pip('uninstall --yes transformers', ignore=True, quiet=True)
pip(f'install tokenizers=={target_tokenizers}', ignore=False, quiet=True)
pip(f'install git+https://github.com/huggingface/transformers@{target_commit}', ignore=False, quiet=True)
@@ -942,6 +947,12 @@ def check_torch():
install(torch_command, 'torch torchvision', quiet=False)
try:
+ try:
+ # import torch pulls torch.distributed immediately which is slow and unnecessary
+ import torch.distributed.tensor._ops as dtensor_ops
+ dtensor_ops.single_dim_strategy._resolve_foreach_elementwise_overload = lambda *a, **kw: None # pylint: disable=protected-access
+ except Exception:
+ pass
import torch
try:
import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
@@ -1355,6 +1366,7 @@ def install_requirements():
# set environment variables controlling the behavior of various libraries
def set_environment():
log.debug('Setting environment tuning')
+ os.environ.setdefault('SDNQ_REGISTER_DIFFUSERS', '1')
os.environ.setdefault('ACCELERATE', 'True')
os.environ.setdefault('ATTN_PRECISION', 'fp16')
os.environ.setdefault('ClDeviceGlobalMemSizeAvailablePercent', '100')
@@ -1365,6 +1377,7 @@ def set_environment():
os.environ.setdefault('CUDA_MODULE_LOADING', 'LAZY')
os.environ.setdefault('DO_NOT_TRACK', '1')
os.environ.setdefault('FORCE_CUDA', '1')
+ os.environ.setdefault('DIFFUSERS_TRUST_REMOTE_KERNELS', 'true')
os.environ.setdefault('GRADIO_ANALYTICS_ENABLED', 'False')
os.environ.setdefault('K_DIFFUSION_USE_COMPILE', '0')
os.environ.setdefault('KINETO_LOG_LEVEL', '3')
diff --git a/launch.py b/launch.py
index 5357cbedb..e3b0b995e 100755
--- a/launch.py
+++ b/launch.py
@@ -196,8 +196,8 @@ def clean_server():
def start_server(immediate=True, server=None):
if args.profile:
import cProfile
- pr = cProfile.Profile()
- pr.enable()
+ profiler = cProfile.Profile()
+ profiler.enable()
import gc
import importlib.util
collected = 0
@@ -221,10 +221,12 @@ def start_server(immediate=True, server=None):
server.wants_restart = False
uvicorn = server.webui(restart=not immediate, _exit=True)
else:
- uvicorn = server.webui(restart=not immediate)
+ uvicorn = server.webui(restart=not immediate, profiler=profiler if args.profile else None)
if args.profile:
- pr.disable()
- installer.print_profile(pr, 'WebUI')
+ profiler.disable()
+ installer.print_profile(profiler, 'WebUI')
+ profiler.clear()
+ profiler.enable()
rec('server')
return uvicorn, server
@@ -252,7 +254,7 @@ def main():
installer.check_version()
installer.check_venv()
log.info(f'Args: {sys.argv[1:]}')
- if not args.skip_env and not args.skip_all:
+ if not args.skip_env:
installer.set_environment()
if args.uv and shutil.which('uv') is None:
installer.install('uv', 'uv')
@@ -343,10 +345,12 @@ def main():
if uv is not None and uv.wants_restart:
clean_server()
log.info('Server restarting...')
- # uv, instance = start_server(immediate=False, server=instance)
os.execv(sys.executable, ['python'] + sys.argv)
else:
log.info('Exiting...')
+ from modules import errors
+ errors.profile_stop()
+ errors.profile_print('Shutdown')
break
time.sleep(1.0)
diff --git a/models/Reference/vladmandic--Mage-Flow-4B.jpg b/models/Reference/vladmandic--Mage-Flow-4B.jpg
new file mode 100644
index 000000000..a90c3c95f
Binary files /dev/null and b/models/Reference/vladmandic--Mage-Flow-4B.jpg differ
diff --git a/models/Reference/vladmandic--Mage-Flow-Turbo-4B.jpg b/models/Reference/vladmandic--Mage-Flow-Turbo-4B.jpg
new file mode 100644
index 000000000..50532d6cf
Binary files /dev/null and b/models/Reference/vladmandic--Mage-Flow-Turbo-4B.jpg differ
diff --git a/modules/api/api.py b/modules/api/api.py
index 108d74fee..3d1bee502 100644
--- a/modules/api/api.py
+++ b/modules/api/api.py
@@ -51,7 +51,6 @@ class Api:
self.add_api_route("/sdapi/v1/status", server.get_status, methods=["GET"], response_model=models.ResStatus, tags=["Server"])
self.add_api_route("/sdapi/v1/platform", server.get_platform, methods=["GET"], tags=["Server"])
self.add_api_route("/sdapi/v1/progress", server.get_progress, methods=["GET"], response_model=models.ResProgress, tags=["Server"])
- self.add_api_route("/sdapi/v1/history", server.get_history, methods=["GET"], response_model=list[models.ResHistory], tags=["Server"])
self.add_api_route("/sdapi/v1/interrupt", server.post_interrupt, methods=["POST"], status_code=204, tags=["Server"])
self.add_api_route("/sdapi/v1/skip", server.post_skip, methods=["POST"], status_code=204, tags=["Server"])
self.add_api_route("/sdapi/v1/shutdown", server.post_shutdown, methods=["POST"], status_code=204, tags=["Server"])
@@ -60,6 +59,8 @@ class Api:
self.add_api_route("/sdapi/v1/cmd-flags", server.get_cmd_flags, methods=["GET"], response_model=models.FlagsModel, tags=["Server"])
self.add_api_route("/sdapi/v1/gpu", gpu.get_gpu, methods=["GET"], tags=["Server"], response_model=list[dict])
self.add_api_route("/sdapi/v1/gpu-smi", gpu.get_gpu_smi, methods=["GET"], response_model=list[models.ResGPU], tags=["Server"])
+ self.add_api_route("/sdapi/v1/history", server.get_history, methods=["GET"], response_model=list[models.ResHistory], tags=["Server"])
+ self.add_api_route("/sdapi/v1/storage", server.get_storage, methods=["GET"], response_model=list[models.ResStorage], tags=["Server"])
# core api using locking
self.add_api_route("/sdapi/v1/txt2img", self.generate.post_text2img, methods=["POST"], response_model=models.ResTxt2Img, tags=["Generation"])
diff --git a/modules/api/middleware.py b/modules/api/middleware.py
index 582369268..6775873a3 100644
--- a/modules/api/middleware.py
+++ b/modules/api/middleware.py
@@ -1,12 +1,14 @@
import ssl
import time
import logging
+import asyncio
from asyncio.exceptions import CancelledError
import anyio
import starlette
import uvicorn
import fastapi
from starlette.responses import JSONResponse
+from starlette.websockets import WebSocket, WebSocketDisconnect
from fastapi import FastAPI, Request, Response
from fastapi.exceptions import HTTPException
from fastapi.responses import RedirectResponse
@@ -30,16 +32,54 @@ def validate_subpath(endpoint: str, subpath: str | None):
return RedirectResponse(url=url, status_code=308)
return None
+class LoopInstrumentorMiddleware:
+ def __init__(self, app):
+ self.app = app
+ self.instrumented = False
+
+ async def __call__(self, scope, receive, send):
+ if not self.instrumented:
+ loop = asyncio.get_running_loop()
+ def verbose_task_factory(loop, coro, context=None):
+ coro_name = getattr(coro, '__qualname__', str(coro))
+ frame = getattr(coro, 'cr_frame', None)
+ origin = f"{frame.f_code.co_filename}:{frame.f_lineno}" if frame else "unknown"
+ log.trace(f"HTTP: coro={coro_name} fn={origin}")
+ if context is not None:
+ return asyncio.Task(coro, loop=loop, name=coro_name, context=context)
+ return asyncio.Task(coro, loop=loop, name=coro_name)
+
+ loop.set_task_factory(verbose_task_factory)
+ self.instrumented = True
+
+ await self.app(scope, receive, send)
+
+
+def setup_logging(debug: bool = False):
+ level = logging.DEBUG if debug else logging.WARNING
+ logging.getLogger("httpcore").setLevel(level)
+ logging.getLogger("httpx").setLevel(level)
+ logging.getLogger("uvicorn.access").setLevel(level)
+ logging.getLogger("asyncio").setLevel(level)
+ if not debug:
+ logging.getLogger("uvicorn.error").disabled = True
+ if debug:
+ asyncio_logger = logging.getLogger("asyncio")
+ if not asyncio_logger.handlers:
+ handler = logging.StreamHandler()
+ handler.setFormatter(logging.Formatter("[asyncio] %(message)s"))
+ asyncio_logger.addHandler(handler)
+
def setup_middleware(app: FastAPI, cmd_opts):
ssl._create_default_https_context = ssl._create_unverified_context # pylint: disable=protected-access
- uvicorn_logger=logging.getLogger("uvicorn.error")
- uvicorn_logger.disabled = True
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
app.middleware_stack = None # reset current middleware to allow modifying user provided list
app.add_middleware(GZipMiddleware, minimum_size=2048)
+ if cmd_opts.profile:
+ app.add_middleware(LoopInstrumentorMiddleware)
if cmd_opts.cors_origins and cmd_opts.cors_regex:
app.add_middleware(CORSMiddleware, allow_origins=cmd_opts.cors_origins.split(','), allow_origin_regex=cmd_opts.cors_regex, allow_methods=['*'], allow_credentials=True, allow_headers=['*'])
elif cmd_opts.cors_origins:
@@ -123,3 +163,20 @@ def setup_middleware(app: FastAPI, cmd_opts):
app.build_middleware_stack() # rebuild middleware stack on-the-fly
log.debug(f'API middleware: {[m.cls.__name__ for m in app.user_middleware]}')
+
+ @app.websocket("/internal/monitor")
+ async def ws_monitor(ws: WebSocket):
+ await ws.accept()
+ try:
+ while True:
+ await asyncio.sleep(1.0)
+ await ws.send_json({"status": "ok"})
+ except WebSocketDisconnect:
+ pass # Expected when client navigates away or closes tab
+ except Exception as e:
+ log.error(f'WebSocket monitor: {e}')
+ finally:
+ try:
+ await ws.close()
+ except RuntimeError:
+ pass # Socket was already closed by client
diff --git a/modules/api/models.py b/modules/api/models.py
index 152a2dbda..7bb0f8f27 100644
--- a/modules/api/models.py
+++ b/modules/api/models.py
@@ -443,7 +443,6 @@ class ReqGetLog(BaseModel):
lines: int = Field(default=100, title="Lines", description="How many lines to return")
clear: bool = Field(default=False, title="Clear", description="Should the log be cleared after returning the lines?")
-
class ReqPostLog(BaseModel):
json: dict | None = Field(default=None, title="Data", description="The data to log")
message: str | None = Field(default=None, title="Message", description="The info message to log")
@@ -453,6 +452,10 @@ class ReqPostLog(BaseModel):
class ReqHistory(BaseModel):
id: int | str | None = Field(default=None, title="Task ID", description="Task ID")
+class ReqStorage(BaseModel):
+ folder: str | None = Field(default=None, title="Folder", description="Storage folder(s)")
+ types: str | None = Field(default=None, title="Types", description="Storage types to filter by")
+
class ReqProgress(BaseModel):
skip_current_image: bool = Field(default=False, title="Skip current image", description="Skip current image serialization")
@@ -472,6 +475,19 @@ class ResHistory(BaseModel):
duration: float | None = Field(default=None, title="Duration", description="Job duration")
outputs: list[str] = Field(title="Outputs", description="List of filenames")
+class ResStorage(BaseModel):
+ name: str = Field(title="Name", description="Storage location name")
+ type: str = Field(title="Type", description="Storage location type")
+ folders: list[str] = Field(title="Folders", description="List of folders in the storage location")
+ paths: list[str] = Field(title="Paths", description="List of resolved paths in the storage location")
+ size: int = Field(title="Size", description="Total size of the storage location in bytes")
+ mtime: float = Field(title="Last modified", description="Last modified timestamp of the storage location")
+ nfiles: int = Field(title="Files", description="Number files in the storage location")
+ nfolders: int = Field(title="Folders", description="Number of folders in the storage location")
+ nsymlinks: int = Field(title="Symlinks", description="Number of symbolic links in the storage location")
+ nerrors: int = Field(title="Errors", description="Number of errors in the storage location")
+ time: float = Field(title="Time", description="Time taken to scan the storage location in seconds")
+
class ResStatus(BaseModel):
status: str = Field(title="Status", description="Current status")
task: str = Field(title="Task", description="Current job")
diff --git a/modules/api/process.py b/modules/api/process.py
index 1a5027714..6264fb741 100644
--- a/modules/api/process.py
+++ b/modules/api/process.py
@@ -1,7 +1,6 @@
from threading import Lock
from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
from fastapi.responses import JSONResponse
-from fastapi.exceptions import HTTPException
from modules.api.helpers import decode_base64_to_image, encode_pil_to_base64
from modules import errors, shared, postprocessing
from modules.api import models, helpers
@@ -20,6 +19,7 @@ class ResPreprocess(BaseModel):
model: str = Field(default='', title="Model", description="The processor model used")
image: str = Field(default='', title="Image", description="The processed image in base64 format")
+
class ReqMask(BaseModel):
image: str = Field(title="Image", description="The base64 encoded image")
type: str = Field(title="Mask type", description="Type of masking image to return")
@@ -27,6 +27,10 @@ class ReqMask(BaseModel):
model: str | None = Field(title="Model", description="The model to use for preprocessing")
params: dict | None = Field(default={}, title="Settings", description="Preprocessor settings")
+class ResMask(BaseModel):
+ mask: str = Field(default='', title="Image", description="The processed image in base64 format")
+
+
class ReqFace(BaseModel):
image: str = Field(title="Image", description="The base64 encoded image")
model: str | None = Field(title="Model", description="The model to use for detection")
@@ -38,8 +42,6 @@ class ResFace(BaseModel):
images: list[str] = Field(title="Image", description="The base64 encoded images of detected faces")
scores: list[float] = Field(title="Scores", description="The scores of the detected faces")
-class ResMask(BaseModel):
- mask: str = Field(default='', title="Image", description="The processed image in base64 format")
class ItemPreprocess(BaseModel):
name: str = Field(title="Name", description="Preprocessor name")
@@ -215,51 +217,37 @@ class APIProcess:
seed = req.seed or -1
seed = processing_helpers.get_fixed_seed(seed)
prompt = ''
- if req.type in ('text', 'image'):
- from modules.scripts_manager import scripts_txt2img
- default_model = 'google/gemma-3-4b-it' if req.type == 'image' else 'google/gemma-3-1b-it'
- model = default_model if req.model is None or len(req.model) < 4 else req.model
- instance = [s for s in scripts_txt2img.scripts if 'prompt_enhance.py' in s.filename][0]
- prompt = instance.enhance(
- model=model,
- prompt=req.prompt,
- system=req.system_prompt,
- prefix=req.prefix,
- suffix=req.suffix,
- sample=req.do_sample,
- min_tokens=req.min_tokens,
- max_tokens=req.max_tokens,
- temperature=req.temperature,
- penalty=req.repetition_penalty,
- top_k=req.top_k,
- top_p=req.top_p,
- thinking=req.thinking,
- keep_thinking=req.keep_thinking,
- 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,
- seed=seed,
- nsfw=req.nsfw,
- custom_args=req.custom_args,
- process_words=req.process_words,
- semantic_threshold=req.semantic_threshold,
- embedding_similarity=req.embedding_similarity,
- use_openai=req.use_openai
- )
- elif req.type == 'video':
- from modules.ui_video_vlm import enhance_prompt
- model = 'Google Gemma 3 4B' if req.model is None or len(req.model) < 4 else req.model
- prompt = enhance_prompt(
- enable=True,
- image=decode_base64_to_image(req.image),
- prompt=req.prompt,
- model=model,
- system_prompt=req.system_prompt,
- nsfw=req.nsfw,
- )
- else:
- raise HTTPException(status_code=400, detail="prompt enhancement: invalid type")
+ from modules.scripts_manager import scripts_txt2img
+ default_model = 'google/gemma-3-4b-it' if req.type == 'image' else 'google/gemma-3-1b-it'
+ model = default_model if req.model is None or len(req.model) < 4 else req.model
+ instance = [s for s in scripts_txt2img.scripts if 'prompt_enhance.py' in s.filename][0]
+ prompt = instance.enhance(
+ model=model,
+ prompt=req.prompt,
+ system=req.system_prompt,
+ prefix=req.prefix,
+ suffix=req.suffix,
+ sample=req.do_sample,
+ min_tokens=req.min_tokens,
+ max_tokens=req.max_tokens,
+ temperature=req.temperature,
+ penalty=req.repetition_penalty,
+ top_k=req.top_k,
+ top_p=req.top_p,
+ thinking=req.thinking,
+ keep_thinking=req.keep_thinking,
+ 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,
+ seed=seed,
+ nsfw=req.nsfw,
+ custom_args=req.custom_args,
+ process_words=req.process_words,
+ semantic_threshold=req.semantic_threshold,
+ embedding_similarity=req.embedding_similarity,
+ use_openai=req.use_openai
+ )
res = models.ResPromptEnhance(prompt=prompt, seed=seed)
return res
diff --git a/modules/api/server.py b/modules/api/server.py
index 65404ca39..79edfb3cc 100644
--- a/modules/api/server.py
+++ b/modules/api/server.py
@@ -111,6 +111,15 @@ def get_history(req: models.ReqHistory = Depends()):
res = [models.ResHistory(**item) for item in res]
return res
+def get_storage(req: models.ReqStorage = Depends()):
+ from modules.storage import check_storage
+ res = check_storage(folders=req.folder,
+ types=req.types.split(',') if req.types else None,
+ silent=True,
+ )
+ res = [models.ResStorage(**loc.dict()) for loc in res]
+ return res
+
def get_progress(req: models.ReqProgress = Depends()):
if shared.state.job_count == 0 and shared.state.sampling_step == 0: # truly idle
return models.ResProgress(id=shared.state.id, progress=0, eta_relative=0, state=shared.state.dict(), textinfo=shared.state.textinfo)
diff --git a/modules/api/validate.py b/modules/api/validate.py
index 5a3755cc3..9738c0efb 100644
--- a/modules/api/validate.py
+++ b/modules/api/validate.py
@@ -41,7 +41,7 @@ log_exclude_prefix = ['/assets']
class Limiter():
- def __init__(self, limit, subpath=None):
+ def __init__(self, limit, subpath=None, debug=False):
import limits
self.request_backend = limits.storage.MemoryStorage()
self.request_limit = limit # default is 300 requests per minute
@@ -53,6 +53,7 @@ class Limiter():
self.log_limiter = limits.parse(f"{self.log_limit}/minute")
self.summary = {}
self.subpath = subpath
+ self.debug = debug
log.info(f'API: limit={self.request_limit} strategy={self.request_strategy.__class__.__name__} backend={self.request_backend.__class__.__name__} subpath={self.subpath}')
@@ -75,6 +76,8 @@ class Limiter():
return status
def check_log(self, client: str, api: str):
+ if self.debug:
+ return True
if self.log_limit < 0:
return True
if any(api.endswith(s) for s in log_exclude_suffix):
@@ -99,7 +102,7 @@ def validate_request(client, endpoint):
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)
+ limiter = Limiter(opts.server_rate_limit, cmd_opts.subpath, cmd_opts.profile)
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 rate limiting
diff --git a/modules/attention.py b/modules/attention.py
index 5adc483ae..7380b5210 100644
--- a/modules/attention.py
+++ b/modules/attention.py
@@ -291,7 +291,7 @@ def set_diffusers_attention(pipe, quiet = False):
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.14.1'
+ 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)
@@ -309,7 +309,7 @@ def get_hf_api_hijack(user_agent = None): # pylint: disable=unused-argument
def hijack_kernels():
global orig_get_kernel # pylint: disable=global-statement
try:
- install('kernels==0.14.1')
+ install('kernels==0.16.0')
import kernels
import kernels.utils
log.debug(f'Attention dispatcher: kernels={kernels.__version__}')
diff --git a/modules/call_queue.py b/modules/call_queue.py
index 71eee8000..88fc5d170 100644
--- a/modules/call_queue.py
+++ b/modules/call_queue.py
@@ -3,7 +3,6 @@ import sys
import html
import threading
import time
-import cProfile
from modules import shared, progress, errors, timer
from modules.logger import log
@@ -65,8 +64,9 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None):
jobid = shared.state.begin(job_name, task_id=task_id)
try:
if shared.cmd_opts.profile:
- pr = cProfile.Profile()
- pr.enable()
+ errors.profile_stop()
+ errors.profile_print('BeforeWrapGradioCall')
+ errors.profile_start()
res = func(*args, **kwargs)
if res is None:
msg = "No result returned from function"
@@ -76,8 +76,9 @@ def wrap_gradio_call(func, extra_outputs=None, add_stats=False, name=None):
else:
res = list(res)
if shared.cmd_opts.profile:
- pr.disable()
- errors.profile(pr, 'Wrap')
+ errors.profile_stop()
+ errors.profile_print('AfterWrapGradioCall')
+ errors.profile_start()
except Exception as e:
errors.display(e, 'gradio call')
res = extra_outputs_array or []
diff --git a/modules/caption/deepbooru.py b/modules/caption/deepbooru.py
index 88095ed01..facd30fd0 100644
--- a/modules/caption/deepbooru.py
+++ b/modules/caption/deepbooru.py
@@ -93,7 +93,7 @@ class DeepDanbooru:
return ''
pic = pil_image.resize((512, 512), resample=Image.Resampling.LANCZOS).convert("RGB")
a = np.expand_dims(np.array(pic, dtype=np.float32), 0) / 255
- with devices.inference_context():
+ with devices.llm_context():
x = torch.from_numpy(a).to(device=devices.device, dtype=devices.dtype)
y = self.model(x)[0].detach().float().cpu().numpy()
probability_dict = {}
diff --git a/modules/caption/joycaption.py b/modules/caption/joycaption.py
index ee527bc70..a96226f73 100644
--- a/modules/caption/joycaption.py
+++ b/modules/caption/joycaption.py
@@ -108,7 +108,7 @@ def predict(question: str, image, vqa_model: str | None = None) -> str:
inputs = processor(text=[convo_string], images=[image], return_tensors="pt").to(devices.device)
inputs['pixel_values'] = inputs['pixel_values'].to(devices.dtype)
try:
- with devices.inference_context():
+ with devices.llm_context():
generate_ids = llava_model.generate( # Generate the captions
**inputs,
# input_ids=inputs['input_ids'],
diff --git a/modules/caption/joytag.py b/modules/caption/joytag.py
index e105eb2fe..d4bc98a0c 100644
--- a/modules/caption/joytag.py
+++ b/modules/caption/joytag.py
@@ -1073,7 +1073,7 @@ def predict(image: Image.Image):
load()
image_tensor = prepare_image(image, model.image_size).unsqueeze(0).to(device=devices.device, dtype=devices.dtype)
try:
- with devices.inference_context():
+ with devices.llm_context():
preds = model({'image': image_tensor})
tag_preds = preds['tags'].sigmoid().cpu()
finally:
diff --git a/modules/caption/moondream3.py b/modules/caption/moondream3.py
index 07866adba..eacbec621 100644
--- a/modules/caption/moondream3.py
+++ b/modules/caption/moondream3.py
@@ -105,7 +105,7 @@ def encode_image(image: Image.Image, cache_key: str | None = None):
model = load_model(loaded)
- with devices.inference_context():
+ with devices.llm_context():
encoded = model.encode_image(image)
if cache_key:
@@ -157,7 +157,7 @@ def query(image: Image.Image, question: str, repo: str, stream: bool = False,
else:
image_input = image
- with devices.inference_context():
+ with devices.llm_context():
response = model.query(
image=image_input,
question=question,
@@ -212,7 +212,7 @@ def caption(image: Image.Image, repo: str, length: str = 'normal', stream: bool
debug(f'LLM: handler=moondream3 method=caption length={length} stream={stream} settings={settings}')
- with devices.inference_context():
+ with devices.llm_context():
response = model.caption(
image,
length=length,
@@ -244,7 +244,7 @@ def point(image: Image.Image, object_name: str, repo: str):
debug(f'LLM: handler=moondream3 method=point object_name="{object_name}"')
- with devices.inference_context():
+ with devices.llm_context():
result = model.point(image, object_name)
debug(f'LLM: handler=moondream3 point_raw_result="{result}" type={type(result)}')
@@ -281,7 +281,7 @@ def detect(image: Image.Image, object_name: str, repo: str, max_objects: int = 1
debug(f'LLM: handler=moondream3 method=detect object_name="{object_name}" max_objects={max_objects}')
- with devices.inference_context():
+ with devices.llm_context():
result = model.detect(image, object_name)
debug(f'LLM: handler=moondream3 detect_raw_result="{result}" type={type(result)}')
diff --git a/modules/caption/vqa.py b/modules/caption/vqa.py
index a2c4bee58..784f6d941 100644
--- a/modules/caption/vqa.py
+++ b/modules/caption/vqa.py
@@ -519,7 +519,7 @@ class VQA:
attention_mask = torch.ones_like(input_ids, device=devices.device)
px = self.model.get_vision_tower().image_processor(images=image, return_tensors="pt")
px = px["pixel_values"].to(self.model.device, dtype=self.model.dtype)
- with devices.inference_context():
+ with devices.llm_context():
outputs = self.model.generate(
inputs=input_ids,
attention_mask=attention_mask,
@@ -673,7 +673,7 @@ class VQA:
defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs}
log.debug(f'LLM: defaults={defaults}')
- with devices.inference_context():
+ with devices.llm_context():
output_ids = self.model.generate(
**inputs,
**gen_kwargs,
@@ -812,7 +812,7 @@ class VQA:
defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs}
log.debug(f'LLM: defaults={defaults}')
- with devices.inference_context():
+ with devices.llm_context():
generation = self.model.generate(
**inputs,
**gen_kwargs,
@@ -899,7 +899,7 @@ class VQA:
defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs}
log.debug(f'LLM: defaults={defaults}')
- with devices.inference_context():
+ with devices.llm_context():
generation = self.model.generate(**inputs, **gen_kwargs)
generation = generation[0][input_len:]
response = self.processor.decode(generation, skip_special_tokens=True)
@@ -932,7 +932,7 @@ class VQA:
question = question.replace('<', '').replace('>', '').replace('_', ' ')
model_inputs = self.processor(text=question, images=image, return_tensors="pt").to(devices.device, devices.dtype)
input_len = model_inputs["input_ids"].shape[-1]
- with devices.inference_context():
+ with devices.llm_context():
generation = self.model.generate(
**model_inputs,
**get_kwargs(self.model),
@@ -989,7 +989,7 @@ class VQA:
if pixel_values is not None:
pixel_values = pixel_values.to(dtype=visual_tokenizer.dtype, device=visual_tokenizer.device)
pixel_values = [pixel_values]
- with devices.inference_context():
+ with devices.llm_context():
output_ids = self.model.generate(
input_ids,
pixel_values=pixel_values,
@@ -1091,7 +1091,7 @@ class VQA:
defaults = {k: v for k, v in helpers.get_default_args(self.model).items() if k not in gen_kwargs}
log.debug(f'LLM: defaults={defaults}')
- with devices.inference_context():
+ with devices.llm_context():
output_ids = self.model.generate(
**inputs,
**gen_kwargs,
@@ -1135,7 +1135,7 @@ class VQA:
input_ids = [self.processor.tokenizer.cls_token_id] + input_ids
input_ids = torch.tensor(input_ids).unsqueeze(0)
git_dict['input_ids'] = input_ids.to(devices.device)
- with devices.inference_context():
+ with devices.llm_context():
generated_ids = self.model.generate(**git_dict)
response = self.processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
@@ -1164,7 +1164,7 @@ class VQA:
move_aux_to_gpu('vqa')
inputs = self.processor(image, question, return_tensors="pt")
inputs = inputs.to(devices.device, devices.dtype)
- with devices.inference_context():
+ with devices.llm_context():
outputs = self.model.generate(**inputs)
response = self.processor.decode(outputs[0], skip_special_tokens=True)
return response
@@ -1193,7 +1193,7 @@ class VQA:
move_aux_to_gpu('vqa')
inputs = self.processor(image, question, return_tensors="pt")
inputs = inputs.to(devices.device)
- with devices.inference_context():
+ with devices.llm_context():
outputs = self.model(**inputs)
logits = outputs.logits
idx = logits.argmax(-1).item()
@@ -1227,7 +1227,7 @@ class VQA:
else:
inputs = self.processor(images=image, return_tensors="pt")
inputs = {k: v.to(devices.device, devices.dtype) if v.is_floating_point() else v.to(devices.device) for k, v in inputs.items()}
- with devices.inference_context():
+ with devices.llm_context():
outputs = self.model.generate(**inputs)
response = self.processor.decode(outputs[0], skip_special_tokens=True)
return response
@@ -1258,7 +1258,7 @@ class VQA:
self._load_moondream(repo)
move_aux_to_gpu('vqa')
question = question.replace('<', '').replace('>', '').replace('_', ' ')
- with devices.inference_context():
+ with devices.llm_context():
if question == 'CAPTION':
response = self.model.caption(image, length="short")['caption']
elif question == 'DETAILED CAPTION':
@@ -1379,7 +1379,7 @@ class VQA:
gen_kwargs['decoder_start_token_id'] = bos_token_id
debug(f'LLM: handler=florence setting decoder_start_token_id={bos_token_id}')
debug(f'LLM: handler=florence generation_kwargs={gen_kwargs}')
- with devices.inference_context(), devices.bypass_sdpa_hijacks():
+ with devices.llm_context():
generated_ids = self.model.generate(
input_ids=input_ids,
pixel_values=pixel_values,
@@ -1431,7 +1431,7 @@ class VQA:
'mask_prompts': None,
'tokenizer': self.processor,
}
- with devices.inference_context():
+ with devices.llm_context():
return_dict = self.model.predict_forward(**input_dict)
response = return_dict["prediction"] # the text format answer
return response
diff --git a/modules/control/proc/marigold/marigold_pipeline.py b/modules/control/proc/marigold/marigold_pipeline.py
index 4ab8c0495..57e6596fd 100644
--- a/modules/control/proc/marigold/marigold_pipeline.py
+++ b/modules/control/proc/marigold/marigold_pipeline.py
@@ -55,7 +55,7 @@ class MarigoldDepthOutput(BaseOutput):
depth_np: np.ndarray
depth_colored: Image.Image
- uncertainty: Union[None, np.ndarray]
+ uncertainty: Union[np.ndarray, None]
class MarigoldPipeline(DiffusionPipeline):
diff --git a/modules/devices.py b/modules/devices.py
index 17dddc1ee..34e96f700 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -242,14 +242,25 @@ def torch_gc(force: bool = False, fast: bool = False, reason: str | None = None)
if force:
# actual gc
collected = gc.collect() if not fast else 0 # python gc
- if cuda_ok:
- try:
- with torch.cuda.device(get_cuda_device_string()):
- torch.cuda.synchronize()
- torch.cuda.empty_cache() # cuda gc
+ try:
+ if hasattr(torch, "accelerator") and torch.accelerator.is_available(): # torch >= 2.6
+ torch.accelerator.synchronize()
+ torch.accelerator.empty_cache()
+ if torch.cuda.is_available() and hasattr(torch.cuda, "ipc_collect"):
torch.cuda.ipc_collect()
- except Exception as e:
- log.error(f'GC: {e}')
+ elif hasattr(torch, "xpu") and hasattr(torch.xpu, "ipc_collect"):
+ torch.xpu.ipc_collect()
+ elif torch.cuda.is_available(): # Fallback for older PyTorch versions
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+ torch.cuda.ipc_collect()
+ elif hasattr(torch, "xpu") and torch.xpu.is_available():
+ torch.xpu.synchronize()
+ torch.xpu.empty_cache()
+ if hasattr(torch.xpu, "ipc_collect"):
+ torch.xpu.ipc_collect()
+ except Exception as e:
+ log.error(f'Torch GC: {e}')
else:
return gpu, ram
t1 = time.time()
@@ -752,3 +763,24 @@ def bypass_sdpa_hijacks():
torch.nn.functional.scaled_dot_product_attention = current_sdpa
if debug:
log.debug('SDPA bypass: restored hijacked attention')
+
+
+@contextlib.contextmanager
+def llm_context():
+ """
+ Combined context manager that applies both inference_context and bypass_sdpa_hijacks.
+ """
+ with inference_context(), bypass_sdpa_hijacks():
+ yield
+
+
+def torch_reset() -> bool:
+ """
+ Resets PyTorch execution graph, flushes VRAM caches, and syncs streams.
+ """
+ torch_gc(force=True, reason='reset', fast=True)
+ try:
+ if hasattr(torch, "_dynamo"):
+ torch._dynamo.reset() # pylint: disable=protected-access
+ except Exception as e:
+ log.error(f"Torch reset: {e}")
diff --git a/modules/errors.py b/modules/errors.py
index 5faf6f9b7..c06d08a4e 100644
--- a/modules/errors.py
+++ b/modules/errors.py
@@ -8,6 +8,7 @@ log = get_log()
setup_logging()
install_traceback()
already_displayed = {}
+_profiler = None
class ValidationError(ValueError):
@@ -63,19 +64,21 @@ def exception(suppress=None):
console.print_exception(show_locals=False, max_frames=16, extra_lines=2, suppress=suppress, theme="ansi_dark", word_wrap=False, width=min([console.width, 200]))
-def profile(profiler, msg: str, n: int = 16):
- profiler.disable()
+def profile_print(msg: str='', n: int = 16, local_profiler=None):
+ if local_profiler is None:
+ local_profiler = _profiler
+ if local_profiler is None:
+ return
import io
import pstats
stream = io.StringIO() # pylint: disable=abstract-class-instantiated
- p = pstats.Stats(profiler, stream=stream)
+ p = pstats.Stats(local_profiler, stream=stream)
p.sort_stats(pstats.SortKey.CUMULATIVE)
p.print_stats(200)
# p.print_title()
# p.print_call_heading(10, 'time')
# p.print_callees(10)
# p.print_callers(10)
- profiler = None
lines = stream.getvalue().split('\n')
lines = [x for x in lines if ' str | None:
+@lru_cache(maxsize=1024)
+def real_path(directory_path: str) -> str | None:
+ """Cached real_path resolution to avoid repeated abspath/expanduser calls."""
+ if not directory_path:
+ return None
try:
return os.path.abspath(os.path.expanduser(directory_path))
except Exception:
- pass
- return None
+ return None
@dataclass(frozen=True)
class Directory(Directory): # pylint: disable=E0102
path: str = field(default_factory=str)
- mtime: float = field(default_factory=float, init=False)
files: FilePathList = field(default_factory=list)
directories: DirectoryPathList = field(default_factory=list)
-
- def __post_init__(self):
- object.__setattr__(self, 'mtime', self.live_mtime)
+ mtime: float = 0.0
@classmethod
def from_dict(cls, dict_object: dict) -> Directory:
directory = cls.__new__(cls)
object.__setattr__(directory, 'path', dict_object.get('path'))
- object.__setattr__(directory, 'mtime', dict_object.get('mtime'))
- object.__setattr__(directory, 'files', dict_object.get('files'))
- object.__setattr__(directory, 'directories', dict_object.get('directories'))
+ object.__setattr__(directory, 'mtime', dict_object.get('mtime', 0.0))
+ object.__setattr__(directory, 'files', dict_object.get('files', []))
+ object.__setattr__(directory, 'directories', dict_object.get('directories', []))
return directory
def clear(self) -> None:
@@ -59,12 +59,15 @@ class Directory(Directory): # pylint: disable=E0102
}))
def update(self, source_directory: Directory) -> Directory:
- if source_directory is not self:
+ if source_directory is not self and source_directory is not None:
self._update(source_directory)
return self
- def _update(self, source:Directory) -> None:
- assert not source.path or source.path == self.path, f'When updating a directory, the paths must match. Attempted to update Directory `{self.path}` with `{source.path}`'
+ def _update(self, source: Directory) -> None:
+ assert not source.path or source.path == self.path, (
+ f'When updating a directory, the paths must match. '
+ f'Attempted to update Directory `{self.path}` with `{source.path}`'
+ )
for dead_path in self.directories:
if dead_path not in source.directories:
delete_cached_directory(dead_path)
@@ -74,105 +77,115 @@ class Directory(Directory): # pylint: disable=E0102
@property
def exists(self) -> bool:
- return self.path and os.path.exists(self.path)
+ return bool(self.path and os.path.exists(self.path))
@property
def is_directory(self) -> bool:
- return self.exists and os.path.isdir(self.path)
+ return bool(self.path and os.path.isdir(self.path))
@property
def live_mtime(self) -> float:
- return os.path.getmtime(self.path) if self.is_directory else 0
+ try:
+ return os.path.getmtime(self.path) if self.path else 0.0
+ except OSError:
+ return 0.0
@property
def is_stale(self) -> bool:
- return not self.is_directory or self.mtime != self.live_mtime
+ return self.mtime != self.live_mtime
class DirectoryCache(UserDict, DirectoryCollection):
def __delattr__(self, directory_path: str) -> None:
directory: Directory = get_directory(directory_path, fetch=False)
if directory:
- map(delete_cached_directory, directory.directories)
+ for child in directory.directories:
+ delete_cached_directory(child)
directory.clear()
- del self.data[directory_path]
+ self.data.pop(directory_path, None)
-def clean_directory(directory: Directory, /, recursive: RecursiveType=False) -> bool:
+def clean_directory(directory: Directory, /, recursive: RecursiveType = False) -> bool:
if not directory.is_directory:
- is_clean = False
delete_cached_directory(directory.path)
- else:
- is_clean = not directory.is_stale
- if not is_clean:
- directory.update(fetch_directory(directory.path))
- else:
- for directory_path in directory.directories[:]:
- try:
- recurse = recursive and (not callable(recursive) or recursive(directory.path))
- directory = get_directory(directory_path, fetch=recurse)
- if directory:
- if directory.is_directory:
- if recurse:
- is_clean = clean_directory(directory, recursive=recurse) and is_clean
- continue
- delete_cached_directory(directory_path)
- # If we had intended to fetch this directory, but didn't, that means it doesn't exist. Purge.
- if recurse:
- directory.directories.remove(directory_path)
- is_clean = False
- except Exception:
- pass
+ return False
+
+ is_clean = not directory.is_stale
+ if not is_clean:
+ fetched = fetch_directory(directory.path)
+ if fetched:
+ directory.update(fetched)
+ elif recursive:
+ for directory_path in list(directory.directories):
+ try:
+ recurse = recursive and (not callable(recursive) or recursive(directory.path))
+ child_dir = get_directory(directory_path, fetch=recurse)
+ if child_dir:
+ if child_dir.is_directory:
+ if recurse:
+ is_clean = clean_directory(child_dir, recursive=recurse) and is_clean
+ continue
+ delete_cached_directory(directory_path)
+ if recurse:
+ directory.directories.remove(directory_path)
+ is_clean = False
+ except Exception:
+ pass
return is_clean
-def get_directory(directory_or_path: str, /, fetch: bool=True) -> Directory | None:
+def get_directory(directory_or_path: str | Directory, /, fetch: bool = True) -> Directory | None:
if isinstance(directory_or_path, Directory):
if directory_or_path.is_directory:
return directory_or_path
- else:
- directory_or_path = directory_or_path.path
- directory_or_path = real_path(directory_or_path)
- if not cache_folders.get(directory_or_path, None):
+ directory_or_path = directory_or_path.path
+
+ resolved_path = real_path(directory_or_path)
+ if not resolved_path:
+ return None
+
+ if resolved_path not in cache_folders:
if fetch:
- directory = fetch_directory(directory_path=directory_or_path)
+ directory = fetch_directory(directory_path=resolved_path)
if directory and do_cache_folders:
- cache_folders[directory_or_path] = directory
+ cache_folders[resolved_path] = directory
return directory
- else:
- clean_directory(cache_folders[directory_or_path])
- return cache_folders[directory_or_path] if directory_or_path in cache_folders else None
+ return None
+
+ cached = cache_folders[resolved_path]
+ clean_directory(cached)
+ return cache_folders.get(resolved_path)
def fetch_directory(directory_path: str) -> Directory | None:
- directory: Directory
for directory in _walk(directory_path, recurse=False):
- return directory # The return is intentional, we get a generator, we only need the one
+ return directory
return None
-def _walk(top, recurse:RecursiveType=True) -> Directory:
- # reimplemented `path.walk()`
+def _walk(top: str, recurse: RecursiveType = True) -> Iterator[Directory]:
nondirs = []
walk_dirs = []
+ top_mtime = 0.0
+
try:
+ top_mtime = os.path.getmtime(top)
scandir_it = os.scandir(top)
except OSError:
return
+
with scandir_it:
- while True:
- try:
- entry = next(scandir_it)
- except StopIteration:
- break
- if not entry.is_dir():
+ for entry in scandir_it:
+ if not entry.is_dir(follow_symlinks=True):
nondirs.append(entry.path)
else:
if entry.is_symlink() and not os.path.exists(entry.path):
log.error(f'Files broken symlink: {entry.path}')
else:
walk_dirs.append(entry.path)
- yield Directory(top, nondirs, walk_dirs)
+
+ yield Directory(path=top, files=nondirs, directories=walk_dirs, mtime=top_mtime)
+
if recurse:
for new_path in walk_dirs:
if callable(recurse) and not recurse(new_path):
@@ -180,13 +193,13 @@ def _walk(top, recurse:RecursiveType=True) -> Directory:
yield from _walk(new_path, recurse=recurse)
-def _cached_walk(top, recurse:RecursiveType=True) -> Directory:
- top = get_directory(top)
- if not top:
+def _cached_walk(top: str, recurse: RecursiveType = True) -> Iterator[Directory]:
+ top_dir = get_directory(top)
+ if not top_dir:
return
- yield top
+ yield top_dir
if recurse:
- for child_directory in top.directories:
+ for child_directory in top_dir.directories:
if os.path.basename(child_directory).startswith('models--'):
continue
if callable(recurse) and not recurse(child_directory):
@@ -194,28 +207,27 @@ def _cached_walk(top, recurse:RecursiveType=True) -> Directory:
yield from _cached_walk(child_directory, recurse=recurse)
-def walk(top, recurse:RecursiveType=True, cached=True) -> Directory:
+def walk(top: str, recurse: RecursiveType = True, cached: bool = True) -> Iterator[Directory]:
yield from _cached_walk(top, recurse=recurse) if cached else _walk(top, recurse=recurse)
-def delete_cached_directory(directory_path:str) -> bool:
- global cache_folders # pylint: disable=W0602
+def delete_cached_directory(directory_path: str) -> bool:
if directory_path in cache_folders:
del cache_folders[directory_path]
+ return True
+ return False
-def is_directory(dir_path:str) -> bool:
- return dir_path and os.path.exists(dir_path) and os.path.isdir(dir_path)
+def is_directory(dir_path: str) -> bool:
+ return bool(dir_path and os.path.isdir(dir_path))
-def directory_mtime(directory_path:str, /, recursive:RecursiveType=True) -> float:
- return float(max(0, *[directory.mtime for directory in get_directories(directory_path, recursive=recursive)]))
+def directory_mtime(directory_path: str, /, recursive: RecursiveType = True) -> float:
+ dirs = get_directories(directory_path, recursive=recursive)
+ return max((d.mtime for d in dirs), default=0.0)
-def unique_directories(directories:DirectoryPathList, /, recursive:RecursiveType=True) -> DirectoryPathIterator:
- '''Ensure no empty, or duplicates'''
- '''If we are going recursive, then directories that are children of other directories are redundant'''
- ''' @todo this is incredibly inneficient. the hit is small, but it is ugly, no? '''
+def unique_directories(directories: DirectoryPathList, /, recursive: RecursiveType = True) -> DirectoryPathIterator:
directories = sorted(unique_paths(directories), reverse=True)
while directories:
directory = directories.pop()
@@ -231,74 +243,81 @@ def unique_directories(directories:DirectoryPathList, /, recursive:RecursiveType
child_directory = directories[-1][len(directory):]
if child_directory:
next_directory = _directory
- if not callable(recursive):
- _remove_directory = next_directory
- else:
- for sub_directory in child_directory.split(os.path.sep):
- next_directory = os.path.join(next_directory, sub_directory)
- if recursive(next_directory):
- _remove_directory = os.path.join(next_directory, '')
- break
+ _remove_directory = None
+ for sub_directory in child_directory.split(os.path.sep):
+ next_directory = os.path.join(next_directory, sub_directory)
+ if recursive(next_directory):
+ _remove_directory = os.path.join(next_directory, '')
+ break
while _remove_directory and directories:
- _d = directories.pop()
if not directories[-1].startswith(_remove_directory):
- del _remove_directory
+ break
+ directories.pop()
-def unique_paths(directory_paths:DirectoryPathList) -> DirectoryPathIterator:
- realpaths = (real_path(directory_path) for directory_path in filter(bool, directory_paths))
- return {real_directory_path: True for real_directory_path in filter(bool, realpaths)}.keys()
+def unique_paths(directory_paths: DirectoryPathList) -> DirectoryPathIterator:
+ seen = set()
+ for path in directory_paths:
+ if path:
+ r = real_path(path)
+ if r and r not in seen:
+ seen.add(r)
+ yield r
-def get_directories(*directory_paths: DirectoryPathList, fetch:bool=True, recursive:RecursiveType=True) -> DirectoryCollection:
- directory_paths = unique_directories(directory_paths, recursive=recursive)
- directories = (get_directory(directory_path, fetch=fetch) for directory_path in directory_paths)
- return filter(bool, directories)
+def get_directories(*directory_paths: DirectoryPathList, fetch: bool = True, recursive: RecursiveType = True) -> DirectoryCollection:
+ dirs = unique_directories(directory_paths, recursive=recursive)
+ return [d for d in (get_directory(p, fetch=fetch) for p in dirs) if d]
-def directory_files(*directories_or_paths: DirectoryPathList | DirectoryList, recursive: RecursiveType=True) -> FilePathIterator:
- return itertools.chain.from_iterable(
- itertools.chain(
- directory_object.files,
- []
- if not recursive
- else itertools.chain.from_iterable(
- directory_files(directory, recursive=recursive)
- for directory
- in filter(
- bool,
- map(get_directory, filter(((bool if recursive else False) if not callable(recursive) else recursive), directory_object.directories))
- )
- )
- )
- for directory_object
- in filter(bool, map(get_directory, directories_or_paths))
- )
+def directory_files(*directories_or_paths: DirectoryPathList | DirectoryList, recursive: RecursiveType = True) -> FilePathIterator:
+ """Iterative directory file gatherer avoiding deeply nested generator recursion."""
+ visited = set()
+ stack = list(directories_or_paths)
+
+ while stack:
+ item = stack.pop()
+ dir_obj = get_directory(item) if not isinstance(item, Directory) else item
+ if not dir_obj or dir_obj.path in visited:
+ continue
+
+ visited.add(dir_obj.path)
+ yield from dir_obj.files
+
+ if recursive:
+ for child_path in dir_obj.directories:
+ if callable(recursive) and not recursive(child_path):
+ continue
+ stack.append(child_path)
-def extension_filter(ext_filter: ExtensionList | None=None, ext_blacklist: ExtensionList | None=None) -> ExtensionFilter:
- if ext_filter:
- ext_filter = [*map(str.upper, ext_filter)]
- if ext_blacklist:
- ext_blacklist = [*map(str.upper, ext_blacklist)]
- def filter_functon(fp:str):
- return (not ext_filter or any(fp.upper().endswith(ew) for ew in ext_filter)) and (not ext_blacklist or not any(fp.upper().endswith(ew) for ew in ext_blacklist))
- return filter_functon
+def extension_filter(ext_filter: ExtensionList | None = None, ext_blacklist: ExtensionList | None = None) -> ExtensionFilter:
+ """Fast C-level tuple.endswith checks."""
+ valid_exts = tuple(ext.lower() if ext.startswith('.') else f'.{ext.lower()}' for ext in ext_filter) if ext_filter else None
+ black_exts = tuple(ext.lower() if ext.startswith('.') else f'.{ext.lower()}' for ext in ext_blacklist) if ext_blacklist else None
+
+ def filter_function(fp: str) -> bool:
+ fp_lower = fp.lower()
+ if valid_exts and not fp_lower.endswith(valid_exts):
+ return False
+ if black_exts and fp_lower.endswith(black_exts):
+ return False
+ return True
+
+ return filter_function
def not_hidden(filepath: str) -> bool:
return not os.path.basename(filepath).startswith('.')
-def filter_files(file_paths: FilePathList, ext_filter: ExtensionList | None=None, ext_blacklist: ExtensionList | None=None) -> FilePathIterator:
+def filter_files(file_paths: FilePathList, ext_filter: ExtensionList | None = None, ext_blacklist: ExtensionList | None = None) -> FilePathIterator:
return filter(extension_filter(ext_filter, ext_blacklist), file_paths)
-def list_files(*directory_paths:DirectoryPathList, ext_filter: ExtensionList | None=None, ext_blacklist: ExtensionList | None=None, recursive:RecursiveType=True) -> FilePathIterator:
- return filter_files(itertools.chain.from_iterable(
- directory_files(directory, recursive=recursive)
- for directory in get_directories(*directory_paths, recursive=recursive)
- ), ext_filter, ext_blacklist)
+def list_files(*directory_paths: DirectoryPathList, ext_filter: ExtensionList | None = None, ext_blacklist: ExtensionList | None = None, recursive: RecursiveType = True) -> FilePathIterator:
+ raw_files = directory_files(*directory_paths, recursive=recursive)
+ return filter_files(raw_files, ext_filter, ext_blacklist)
cache_folders = DirectoryCache({})
diff --git a/modules/framepack/framepack_ui.py b/modules/framepack/framepack_ui.py
index 92029085a..dbc99e853 100644
--- a/modules/framepack/framepack_ui.py
+++ b/modules/framepack/framepack_ui.py
@@ -12,7 +12,7 @@ def change_sections(duration, mp4_fps, mp4_interpolate, latent_ws, variant):
return gr.update(value=f'Target video: {num_frames} frames in {num_sections} sections'), gr.update(lines=max(2, 2*num_sections//3))
-def create_ui(prompt, negative, styles, _overrides, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
+def create_ui(prompt, negative, styles, _overrides, script_inputs, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
with gr.Row():
with gr.Column(variant='compact', elem_id="framepack_settings", elem_classes=['settings-column'], scale=1):
with gr.Row():
@@ -114,7 +114,7 @@ def create_ui(prompt, negative, styles, _overrides, mp4_fps, mp4_interpolate, mp
framepack_dict = dict(
fn=run_framepack,
_js="submit_framepack",
- inputs=state_inputs + framepack_inputs,
+ inputs=state_inputs + framepack_inputs + script_inputs,
outputs=framepack_outputs,
show_progress='hidden',
)
diff --git a/modules/framepack/framepack_wrappers.py b/modules/framepack/framepack_wrappers.py
index 0e3c546ad..6147d349a 100644
--- a/modules/framepack/framepack_wrappers.py
+++ b/modules/framepack/framepack_wrappers.py
@@ -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, attention, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt):
+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, attention, vae_type, variant, vlm_enhance, vlm_model, vlm_system_prompt, *_args, **_kwargs):
variant = variant or 'bi-directional'
if init_image is None:
init_image = np.zeros((resolution, resolution, 3), dtype=np.uint8)
diff --git a/modules/loader.py b/modules/loader.py
index ed47ab628..99d152491 100644
--- a/modules/loader.py
+++ b/modules/loader.py
@@ -78,6 +78,9 @@ except Exception:
pass
try:
+ # del torch._C._c10d_init
+ import torch.distributed # pylint: disable=ungrouped-imports
+ torch.distributed.is_available = lambda: False
import torch.distributed.distributed_c10d as _c10d # pylint: disable=unused-import,ungrouped-imports
except Exception:
log.warning('Loader: torch is not built with distributed support')
@@ -98,8 +101,9 @@ urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
warnings.filterwarnings(action="ignore", category=UserWarning, module="torchvision")
torchvision = None
try:
+ sys.modules['torchvision.samples'] = types.ModuleType("torchvision.samples") # monkey-patch to avoid torchvision sample download
import torchvision # pylint: disable=W0611,C0411
- import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411
+ # import pytorch_lightning # pytorch_lightning should be imported after torch, but it re-enables warnings on import so import once to disable them # pylint: disable=W0611,C0411
except Exception as e:
report(f'torchvision=={torchvision.__version__ if torchvision is not None else None}', e)
diff --git a/modules/logger.py b/modules/logger.py
index cc0d4f27f..740c74782 100644
--- a/modules/logger.py
+++ b/modules/logger.py
@@ -281,11 +281,7 @@ def setup_logging(debug=None, trace=None, filename=None):
logging.getLogger("lycoris").handlers = log.handlers
logging.getLogger("ControlNet").handlers = log.handlers
- logging.getLogger("asyncio").setLevel(logging.ERROR)
logging.getLogger("diffusers").setLevel(logging.ERROR)
logging.getLogger("transformers").setLevel(logging.ERROR)
- logging.getLogger("httpcore").setLevel(logging.ERROR)
- logging.getLogger("httpx").setLevel(logging.ERROR)
logging.getLogger("torch").setLevel(logging.ERROR)
logging.getLogger("urllib3").setLevel(logging.ERROR)
- logging.getLogger("uvicorn.access").setLevel(logging.ERROR)
diff --git a/modules/lora/lora_apply.py b/modules/lora/lora_apply.py
index b37292f27..54fd86dad 100644
--- a/modules/lora/lora_apply.py
+++ b/modules/lora/lora_apply.py
@@ -139,7 +139,7 @@ def network_calc_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.Grou
return batch_updown, batch_ex_bias
-def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, model_weights: None | torch.Tensor = None, lora_weights: torch.Tensor = None, deactivate: bool = False, device: torch.device = None, bias: bool = False):
+def network_add_weights(self: torch.nn.Conv2d | torch.nn.Linear | torch.nn.GroupNorm | torch.nn.LayerNorm | diffusers.models.lora.LoRACompatibleLinear | diffusers.models.lora.LoRACompatibleConv, model_weights: torch.Tensor | None = None, lora_weights: torch.Tensor = None, deactivate: bool = False, device: torch.device = None, bias: bool = False):
if lora_weights is None:
return
if deactivate:
diff --git a/modules/ltx/ltx_process.py b/modules/ltx/ltx_process.py
index e7bb3ce2a..8c4634407 100644
--- a/modules/ltx/ltx_process.py
+++ b/modules/ltx/ltx_process.py
@@ -3,7 +3,7 @@ import time
import torch
from PIL import Image
-from modules import shared, errors, timer, memstats, progress, processing, sd_models, sd_samplers, devices, extra_networks, call_queue
+from modules import shared, errors, timer, memstats, progress, processing, sd_models, sd_samplers, devices, extra_networks, call_queue, scripts_manager
from modules.logger import log
from modules.ltx import ltx_capabilities
from modules.ltx.ltx_diffusers_patch import apply_patch as apply_ltx_diffusers_patch
@@ -151,6 +151,8 @@ def run_ltx(task_id,
mp4_thumb: bool,
audio_enable: bool,
_overrides,
+ *args,
+ **_kwargs,
):
def abort(e, ok: bool = False, p=None):
@@ -282,13 +284,15 @@ def run_ltx(task_id,
vae_tile_frames=16,
)
processing.fix_seed(p)
- p.scripts = None
- p.script_args = None
p.do_not_save_grid = True
p.do_not_save_samples = not mp4_frames
p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_video)
p.ops.append('video')
+ p.scripts = scripts_manager.scripts_video
+ p.script_args = args
+ processed: processing.Processed = scripts_manager.scripts_video.run(p, *args)
+
p.task_args['num_inference_steps'] = p.steps
p.task_args['width'] = p.width
p.task_args['height'] = p.height
@@ -337,7 +341,9 @@ def run_ltx(task_id,
try:
if needs_latent_path:
- prompt_final, negative_final, networks = get_prompts(prompt, negative, styles)
+ if p.scripts is not None and isinstance(p.scripts, scripts_manager.ScriptRunner):
+ p.scripts.before_process(p)
+ prompt_final, negative_final, networks = get_prompts(p)
extra_networks.activate(p, networks)
# Encode once and reuse across stages; encode_prompt short-circuits when
# embeds are passed to __call__. CPU park keeps them off GPU between stages.
diff --git a/modules/ltx/ltx_ui.py b/modules/ltx/ltx_ui.py
index d15c4aa78..bb326ad68 100644
--- a/modules/ltx/ltx_ui.py
+++ b/modules/ltx/ltx_ui.py
@@ -55,7 +55,7 @@ def _model_change(model_name: str):
)
-def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
+def create_ui(prompt, negative, styles, overrides, script_inputs, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
with gr.Row():
with gr.Column(variant='compact', elem_id="ltx_settings", elem_classes=['settings-column'], scale=1):
with gr.Row():
@@ -67,8 +67,8 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
with gr.Accordion(open=False, label='Size', elem_id='ltx_size_accordion'):
width, height = ui_sections.create_resolution_inputs('ltx', default_width=832, default_height=480)
with gr.Row():
- frames = gr.Slider(label='LTX Frames', minimum=1, maximum=1024, step=1, value=121, elem_id='ltx_frames')
- seed = gr.Number(label='Initial seed', value=-1, elem_id='ltx_seed', container=True)
+ frames = gr.Slider(label='LTX frames', minimum=1, maximum=1024, step=1, value=121, elem_id='ltx_frames')
+ seed = gr.Number(label='LTX seed', value=-1, elem_id='ltx_seed', container=True)
random_seed = ToolButton(ui_symbols.random, elem_id='ltx_seed_random')
random_seed.click(fn=lambda: -1, show_progress='hidden', inputs=[], outputs=[seed])
input_media_accordion = gr.Accordion(open=False, label="Input media", elem_id='ltx_input_media_accordion', visible=False)
@@ -88,16 +88,16 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
condition_video_skip = gr.Slider(label='LTX frames skip', minimum=0, maximum=1024, step=1, value=0, elem_id="ltx_condition_video_sip")
with gr.Tab('Gallery prefix', id='ltx_condition_batch_tab'):
condition_files = gr.Files(label="Image Batch", interactive=True, elem_id="ltx_condition_batch")
- upsample_accordion = gr.Accordion(open=False, label="Upsample", elem_id='ltx_upsample_accordion')
+ upsample_accordion = gr.Accordion(open=False, label="Upscale", elem_id='ltx_upsample_accordion')
with upsample_accordion:
with gr.Row():
- upsample_enable = gr.Checkbox(label='LTX enable upsampling', value=False, elem_id="ltx_upsample_enable")
- upsample_ratio = gr.Slider(label='LTX upsample ratio', minimum=1.0, maximum=4.0, step=0.1, value=2.0, elem_id="ltx_upsample_ratio")
+ upsample_enable = gr.Checkbox(label='LTX upscale', value=False, elem_id="ltx_upsample_enable")
+ upsample_ratio = gr.Slider(label='LTX scale', minimum=1.0, maximum=4.0, step=0.1, value=2.0, elem_id="ltx_upsample_ratio")
refine_accordion = gr.Accordion(open=False, label="Refine", elem_id='ltx_refine_accordion')
with refine_accordion:
with gr.Row():
- refine_enable = gr.Checkbox(label='LTX enable refine', value=False, elem_id="ltx_refine_enable")
- refine_strength = gr.Slider(label='LTX refine strength', minimum=0.1, maximum=1.0, step=0.05, value=0.4, elem_id="ltx_refine_strength")
+ refine_enable = gr.Checkbox(label='LTX refine', value=False, elem_id="ltx_refine_enable")
+ refine_strength = gr.Slider(label='LTX strength', minimum=0.1, maximum=1.0, step=0.05, value=0.4, elem_id="ltx_refine_strength")
parameters_accordion = gr.Accordion(open=False, label="Advanced", elem_id='ltx_parameters_accordion')
with parameters_accordion:
steps, sampler_index = ui_sections.create_sampler_and_steps_selection(None, "ltx", default_steps=50)
@@ -108,7 +108,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
dynamic_shift = gr.Checkbox(label='LTX dynamic shift', value=False, elem_id="ltx_dynamic_shift")
with gr.Row():
decode_timestep = gr.Slider(label='LTX decode timestep', minimum=0.0, maximum=1.0, step=0.01, value=0.05, elem_id="ltx_decode_timestep")
- image_cond_noise_scale = gr.Slider(label='LTX image cond noise scale', minimum=0.0, maximum=1.0, step=0.005, value=0.025, elem_id="ltx_image_cond_noise_scale")
+ image_cond_noise_scale = gr.Slider(label='LTX image cond', minimum=0.0, maximum=1.0, step=0.005, value=0.025, elem_id="ltx_image_cond_noise_scale")
audio_accordion = gr.Accordion(open=False, label="Audio", elem_id='ltx_audio_accordion', visible=False)
with audio_accordion:
with gr.Row():
@@ -175,7 +175,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
video_dict = dict(
fn=ltx_process.run_ltx,
_js="submit_ltx",
- inputs=state_inputs + video_inputs,
+ inputs=state_inputs + video_inputs + script_inputs,
outputs=video_outputs,
show_progress='hidden',
)
diff --git a/modules/ltx/ltx_util.py b/modules/ltx/ltx_util.py
index 032bc04b8..11afb0b1c 100644
--- a/modules/ltx/ltx_util.py
+++ b/modules/ltx/ltx_util.py
@@ -123,7 +123,7 @@ def ltx_scheduler_opts(sd_model, *, dynamic_shift=None, sampler_shift=None):
if orig_flow_shift is not None and hasattr(sd_model.scheduler.config, 'flow_shift'):
sd_model.scheduler.config.flow_shift = orig_flow_shift
sd_model.scheduler.register_to_config(flow_shift=orig_flow_shift)
- log.debug(f'LTX: scheduler/opts restored dynamic_shift={orig_dynamic_shift} sampler_shift={orig_sampler_shift}')
+ # log.debug(f'LTX: scheduler/opts restored dynamic_shift={orig_dynamic_shift} sampler_shift={orig_sampler_shift}')
def _condition_cls(family: str):
@@ -201,9 +201,9 @@ def get_conditions(width, height, condition_strength, condition_images, conditio
return conditions
-def get_prompts(prompt, negative, styles):
- prompt = shared.prompt_styles.apply_styles_to_prompt(prompt, styles)
- negative = shared.prompt_styles.apply_negative_styles_to_prompt(negative, styles)
+def get_prompts(p):
+ prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
+ negative = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
prompts, networks = extra_networks.parse_prompts([prompt])
prompt = prompts[0] if len(prompts) > 0 else prompt
return prompt, negative, networks
diff --git a/modules/memmon.py b/modules/memmon.py
index 32070ac79..4155dedd0 100644
--- a/modules/memmon.py
+++ b/modules/memmon.py
@@ -65,8 +65,8 @@ class MemUsageMonitor:
retries = mem_mon_read.pop("retries")
vram = {k: v//1048576 for k, v in mem_mon_read.items()}
if 'active_peak' in vram:
- peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
- used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+ peak = max(vram.get('active_peak', 0), vram.get('reserved_peak', 0), vram.get('used', 0))
+ used = round(100.0 * peak / vram.get('total', 0)) if vram.get('total', 0) > 0 else 0
else:
peak = 0
used = 0
diff --git a/modules/model_quant.py b/modules/model_quant.py
index e4e528ec4..bd990bcc1 100644
--- a/modules/model_quant.py
+++ b/modules/model_quant.py
@@ -94,7 +94,14 @@ def get_sdnq_devices(mode="pre"):
return quantization_device, return_device
-def create_sdnq_config(kwargs = None, allow: bool = True, module: str = 'Model', weights_dtype: str | None = None, quantized_matmul_dtype: str | None = None, modules_to_not_convert: list | None = None, modules_dtype_dict: dict | None = None):
+def create_sdnq_config(kwargs = None,
+ allow: bool = True,
+ module: str = 'Model',
+ weights_dtype: str | None = None,
+ quantized_matmul_dtype: str | None = None,
+ modules_to_not_convert: list | None = None,
+ modules_dtype_dict: dict | None = None,
+ ):
from modules import shared
if allow and (shared.opts.sdnq_quantize_mode in {'pre', 'auto'}) and (module == 'any' or module in shared.opts.sdnq_quantize_weights):
from modules.sdnq import SDNQConfig
@@ -191,6 +198,8 @@ def check_nunchaku(module: str = ''):
from modules import shared
if 'nunchaku' not in shared.opts.sd_model_checkpoint.lower():
return False
+ if 'nunchaku-lite' in shared.opts.sd_model_checkpoint.lower():
+ return False
base_path = shared.opts.sd_model_checkpoint.split('+')[0]
for v in shared.reference_models.values():
if v.get('path', '') != base_path:
@@ -216,7 +225,12 @@ def create_config(kwargs = None, allow: bool = True, module: str = 'Model', modu
kwargs = {}
if module == 'Model' and dont_quant():
return kwargs
- kwargs = create_sdnq_config(kwargs, allow=allow, module=module, modules_to_not_convert=modules_to_not_convert, modules_dtype_dict=modules_dtype_dict)
+ kwargs = create_sdnq_config(kwargs,
+ allow=allow,
+ module=module,
+ modules_to_not_convert=modules_to_not_convert,
+ modules_dtype_dict=modules_dtype_dict
+ )
if kwargs is not None and 'quantization_config' in kwargs:
if debug:
log.trace(f'Quantization: type=sdnq config={kwargs.get("quantization_config", None)}')
diff --git a/modules/modeldata.py b/modules/modeldata.py
index 032f50a5c..439123a05 100644
--- a/modules/modeldata.py
+++ b/modules/modeldata.py
@@ -139,6 +139,8 @@ def get_model_type(pipe):
model_type = 'sdxs'
elif 'SeFi' in name:
model_type = 'sefi'
+ elif 'Mage-Flow' in name:
+ model_type = 'mageflow'
# video models
elif "Kandinsky5" in name and '2V' in name:
model_type = 'kandinsky5video'
diff --git a/modules/modelloader.py b/modules/modelloader.py
index 282dd1310..1212277ed 100644
--- a/modules/modelloader.py
+++ b/modules/modelloader.py
@@ -407,9 +407,9 @@ def cleanup_models():
def move_files(src_path: str, dest_path: str, ext_filter: str | None = None):
try:
- if not os.path.exists(dest_path):
- os.makedirs(dest_path)
if os.path.exists(src_path):
+ if not os.path.exists(dest_path):
+ os.makedirs(dest_path)
for file in os.listdir(src_path):
fullpath = os.path.join(src_path, file)
if os.path.isfile(fullpath):
diff --git a/modules/modelstats.py b/modules/modelstats.py
index 1d086ce9f..f6eaf9c44 100644
--- a/modules/modelstats.py
+++ b/modules/modelstats.py
@@ -13,43 +13,61 @@ def walk(folder: str):
return files
-def stat(folder: str) -> tuple[int, datetime]:
+def stat(folder: str, follow: bool = False, extended: bool = False, exclude: list[str] = []):
+ _files = 0
+ _folders = 0
+ _symlinks = 0
+ _errors = 0
_size = 0
_mtime = 0.0
def recurse(folder: str):
- nonlocal _size, _mtime
+ nonlocal _size, _mtime, _files, _folders, _symlinks, _errors
with os.scandir(folder) as entries:
for entry in entries:
try:
- if entry.is_file(follow_symlinks=False):
+ if any(part == ex for part in entry.path.split(os.sep) for ex in exclude):
+ continue
+ if entry.is_file(follow_symlinks=follow):
try:
- _stat = entry.stat(follow_symlinks=False)
+ _stat = entry.stat(follow_symlinks=follow)
except Exception:
- _stat = os.stat(entry.path, follow_symlinks=False)
+ _stat = os.stat(entry.path, follow_symlinks=follow)
_size += _stat.st_size
+ _files += 1
if _stat.st_mtime > _mtime:
_mtime = _stat.st_mtime
- elif entry.is_dir(follow_symlinks=False):
+ elif entry.is_symlink():
+ _symlinks += 1
+ elif entry.is_dir(follow_symlinks=follow):
+ _folders += 1
recurse(entry.path)
except (FileNotFoundError, PermissionError):
+ _errors += 1
continue
try:
- if os.path.isfile(folder):
- _stat = os.stat(folder, follow_symlinks=False)
+ s_folder = str(folder)
+ if any(s_folder in ex for ex in exclude):
+ return _size, datetime.fromtimestamp(_mtime).replace(microsecond=0), _files, _folders, _symlinks, _errors
+ elif os.path.isfile(folder):
+ _stat = os.stat(folder, follow_symlinks=follow)
_size = _stat.st_size
_mtime = _stat.st_mtime
+ _files = 1
elif os.path.isdir(folder):
+ _folders = 1
recurse(folder)
else:
pass
except (FileNotFoundError, PermissionError):
- pass
+ _errors += 1
try:
_datetime = datetime.fromtimestamp(_mtime).replace(microsecond=0)
except (OSError, ValueError):
_datetime = datetime.fromtimestamp(0)
+ if extended:
+ return _size, _datetime, _files, _folders, _symlinks, _errors
return _size, _datetime
diff --git a/modules/postprocess/seedvr_model.py b/modules/postprocess/seedvr_model.py
index 5da179635..27f8e775b 100644
--- a/modules/postprocess/seedvr_model.py
+++ b/modules/postprocess/seedvr_model.py
@@ -170,7 +170,7 @@ class UpscalerSeedVR(Upscaler):
devices.torch_gc(fast=True)
t0 = time.time()
with devices.inference_context():
- self.pbar.update(self.task, description=f'inference: batch={self.step}')
+ self.pbar.update(self.task, description='inference')
result = generation.generation_step_original(*args, **kwargs)
self.pbar.update(self.task, advance=self.step)
self.timer.ts('step', t0)
@@ -187,13 +187,38 @@ class UpscalerSeedVR(Upscaler):
image = Image.open(image)
image = image.convert("RGB")
width = image.width
+ height = image.height
tensor = np.array(image)
tensor = torch.from_numpy(tensor).to(device=devices.device, dtype=devices.dtype).unsqueeze(0) / 255.0
self.frames = 1
- return tensor, width
+ return tensor, width, height
except Exception as e:
log.error(f'Upscaler: name="SeedVR2" image="{image}" {e}')
- return None, None
+ return None, None, None
+
+ def read_audio(self, video_path: str):
+ audio_frames = []
+ audio_meta = None
+ try:
+ from modules.video_models.video_utils import check_av
+ av = check_av()
+ container = av.open(video_path)
+ if container.streams.audio:
+ audio_stream = container.streams.audio[0]
+ audio_meta = {
+ "sr": audio_stream.codec_context.sample_rate,
+ "channels": audio_stream.codec_context.channels,
+ "layout": audio_stream.layout.name if audio_stream.layout else "stereo",
+ "format": audio_stream.codec_context.format.name,
+ }
+ for frame in container.decode(audio_stream):
+ audio_frames.append(frame)
+ container.close()
+ except Exception as e:
+ log.error(f'Upscaler: name="SeedVR2" video="{video_path}" {e}')
+ if audio_meta and len(audio_frames) > 0:
+ return {"frames": audio_frames, **audio_meta}
+ return None
def read_video(self, video_path: str):
try:
@@ -201,9 +226,10 @@ class UpscalerSeedVR(Upscaler):
cap = cv2.VideoCapture(video_path)
if not cap.isOpened():
log.error(f'Upscaler: name="SeedVR2" video="{video_path}" failed to open')
- return None, None
+ return None, None, None
frames = []
width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+ height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
self.fps = int(cap.get(cv2.CAP_PROP_FPS))
while True:
ret, frame = cap.read()
@@ -214,20 +240,21 @@ class UpscalerSeedVR(Upscaler):
cap.release()
if len(frames) == 0:
log.error(f'Upscaler: name="SeedVR2" video="{video_path}" no frames read')
- return None, None
+ return None, None, None
tensor = torch.from_numpy(np.array(frames)).to(device=devices.device, dtype=devices.dtype) / 255.0
self.frames = tensor.shape[0]
- return tensor, width
+ return tensor, width, height
except Exception as e:
log.error(f'Upscaler: name="SeedVR2" video="{video_path}" {e}')
- return None, None
+ return None, None, None
- def create_video(self, tensor: torch.Tensor, codec: str = 'libx264', codec_opt: str = 'crf:16', interpolate: int = 0):
+ def create_video(self, tensor: torch.Tensor, audio, codec: str = 'libx264', codec_opt: str = 'crf:16', interpolate: int = 0):
t0 = time.time()
from modules.video_models.video_save import save_video
pixels = tensor.permute(3, 0, 1, 2).unsqueeze(0) # from (t, h, w, c) to (n, c, t, h, w)
_frames, filename, _thumb = save_video(p=None,
pixels=pixels,
+ audio=audio,
mp4_fps=self.fps,
mp4_thumb=False,
mp4_frames=False,
@@ -255,7 +282,7 @@ class UpscalerSeedVR(Upscaler):
interpolate: int = 1,
codec: str = 'libx264',
codec_opt: str = 'crf:16',
- vae_memory: float = 0.2,
+ vae_memory: float = 0.5,
vae_tile_encode: bool = True,
vae_tile_decode: bool = True,
):
@@ -273,11 +300,13 @@ class UpscalerSeedVR(Upscaler):
from modules.seedvr.src.core import generation
+ audio = None
self.scale = self.scale if scale is None else scale
if isinstance(img, Image.Image):
- tensor, width = self.read_image(img)
+ tensor, width, height = self.read_image(img)
elif isinstance(img, str):
- tensor, width = self.read_video(img)
+ tensor, width, height = self.read_video(img)
+ audio = self.read_audio(img)
else:
log.error(f'Upscaler: name="SeedVR2" image="{img}" unsupported type {type(img)}')
return img
@@ -287,6 +316,7 @@ class UpscalerSeedVR(Upscaler):
log.error(f'Upscaler: name="SeedVR2" image="{img}" failed to read')
return img
width = int(self.scale * width) // 8 * 8
+ height = int(self.scale * height) // 8 * 8
random.seed()
seed = int(random.randrange(4294967294)) if seed == -1 else int(seed)
self.step = 1 if self.frames == 1 else batch_size - batch_overlap
@@ -294,7 +324,7 @@ class UpscalerSeedVR(Upscaler):
mode = "mode=image" if self.frames == 1 else f"mode=video frames={self.frames}"
batch_info = f'batch=(size={batch_size} overlap={batch_overlap})'
vae_info = f'vae=(tiled={vae_tile_encode}/{vae_tile_decode} memory={vae_memory} size={tile_size} overlap={tile_overlap})'
- log.info(f'Upscaler: type="{self.name}" model="{selected_file}" {mode} scale={self.scale} cfg={cfg_scale}:{cfg_rescale} seed={seed} steps={steps} offload={self.offload} {batch_info} {vae_info}')
+ log.info(f'Upscaler: type="{self.name}" model="{selected_file}" {mode} scale={self.scale} width={width} height={height} cfg={cfg_scale}:{cfg_rescale} seed={seed} steps={steps} offload={self.offload} {batch_info} {vae_info}')
import rich.progress as rp
self.pbar = rp.Progress(rp.TextColumn('[cyan]SeedVR:'), rp.BarColumn(), rp.MofNCompleteColumn(), rp.TaskProgressColumn(), rp.TimeRemainingColumn(), rp.TimeElapsedColumn(), rp.TextColumn('[cyan]{task.description}'), console=console)
@@ -342,7 +372,7 @@ class UpscalerSeedVR(Upscaler):
if self.frames == 1:
result = convert.to_pil(result_tensor.squeeze())
elif self.frames > 1:
- result = self.create_video(result_tensor, codec=codec, codec_opt=codec_opt, interpolate=interpolate)
+ result = self.create_video(result_tensor, audio, codec=codec, codec_opt=codec_opt, interpolate=interpolate)
else:
log.error(f'Upscaler: name="SeedVR2" model="{selected_file}" no frames generated')
result = img
diff --git a/modules/processing.py b/modules/processing.py
index 9940f77a0..fafd4d9f6 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -423,13 +423,13 @@ def print_stats():
if shared.opts.sdnq_dequantize_compile:
from modules.timer_sdnq import update_sdnq_attention_timers
update_sdnq_attention_timers()
- if timer.autotune.get_total() > 0.001:
+ if timer.autotune.get_total() > 0.1:
log.debug(f'Processed: autotune={timer.autotune.dct(min_time=0)}')
if devices.triton_ok:
from modules.sd_models_compile import update_compile_times
update_compile_times()
- if timer.dynamo.get_total() > 0.001:
+ if timer.dynamo.get_total() > 0.1:
log.debug(f'Processed: dynamo={timer.dynamo.dct(min_time=2.0, no_total=True)}')
diff --git a/modules/processing_callbacks.py b/modules/processing_callbacks.py
index be9d8727a..4fa2eebc3 100644
--- a/modules/processing_callbacks.py
+++ b/modules/processing_callbacks.py
@@ -61,6 +61,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
torch.xpu.synchronize(devices.device)
elif devices.backend in {"cuda", "zluda", "rocm"}:
torch.cuda.synchronize(devices.device)
+ time.sleep(0.001) # 1ms yield frees GIL for the preview thread
t1 = time.time()
@@ -82,7 +83,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
latents = kwargs.get('latents', None)
if debug:
- debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} kwargs={list(kwargs)}')
+ debug_callback(f'Callback: step={step} timestep={timestep} latents={latents.shape if latents is not None else None} sync={shared.opts.torch_sync} kwargs={list(kwargs)}')
if shared.state.sampling_steps == 0 and getattr(pipe, 'num_timesteps', 0) > 0:
shared.state.sampling_steps = pipe.num_timesteps
shared.state.step()
@@ -91,14 +92,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
if latents is None or p is None:
return kwargs
- """
- if torch.isnan(latents).any().item():
- log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="contains NaN values"')
- if (shared.state.current_latent is not None) and (shared.state.current_latent.shape == latents.shape):
- log.error(f'Callback: step={step} timestep={timestep} latents={latents.shape}:{latents.device}:{latents.dtype} error="replacing with previous latent"')
- latents = shared.state.current_latent
- """
-
if len(getattr(p, 'ip_adapter_names', [])) > 0 and p.ip_adapter_names[0] != 'None':
ip_adapter_scales = list(p.ip_adapter_scales)
ip_adapter_starts = list(p.ip_adapter_starts)
@@ -122,6 +115,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
cfg_end = getattr(p, "cfg_end", 1.0) or 1.0
total_steps = getattr(pipe, "num_timesteps", 0)
target_step = int(total_steps * cfg_end) if total_steps else 0
+
if (cfg_end < 1.0) and not getattr(pipe, "_cfg_end_applied", False) and (step >= target_step):
pipe._cfg_end_applied = True # pylint: disable=protected-access
if "PAG" in shared.sd_model.__class__.__name__:
@@ -145,7 +139,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
else:
width = getattr(p, 'width', 1024)
height = getattr(p, 'height', 1024)
- shared.state.current_latent = pipe._unpack_latents(kwargs['latents'], height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
+ shared.state.current_latent = pipe._unpack_latents(latents, height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
if current_noise_pred is not None:
shared.state.current_noise_pred = pipe._unpack_latents(current_noise_pred, height, width, pipe.vae_scale_factor) # pylint: disable=protected-access
else:
@@ -158,7 +152,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
else:
width = getattr(p, 'width', 1024)
height = getattr(p, 'height', 1024)
- latents = kwargs['latents']
if len(latents.shape) == 4:
latents = pipe._unpatchify_latents(latents) # [B, C*4, h/2, w/2] -> [B, C, h, w] # pylint: disable=protected-access
elif len(latents.shape) == 3: # packed format [B, seq_len, patch_channels]
@@ -183,7 +176,6 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
current_noise_pred = current_noise_pred.permute(0, 3, 1, 4, 2, 5).reshape(b, channels, h_patches * 2, w_patches * 2)
shared.state.current_noise_pred = current_noise_pred
elif 'Ideogram4' in pipe.__class__.__name__: # packed normalized [B, seq, 128] -> Flux.2 latent space for TAE FLUX.2
- latents = kwargs['latents']
if latents.ndim == 3:
b, seq_len, packed_ch = latents.shape
vae_scale = getattr(pipe, 'vae_scale_factor', 8)
@@ -203,7 +195,7 @@ def diffusers_callback(pipe, step: int = 0, timestep: int = 0, kwargs: dict | No
shared.state.current_latent = latents
shared.state.current_noise_pred = current_noise_pred
else:
- shared.state.current_latent = kwargs['latents']
+ shared.state.current_latent = latents
shared.state.current_noise_pred = current_noise_pred
# Video latent preview: extract middle frame from 5D [B,C,T,H,W] to 4D [B,C,H,W]
diff --git a/modules/progress.py b/modules/progress.py
index 5c2c3c7ec..890ead043 100644
--- a/modules/progress.py
+++ b/modules/progress.py
@@ -48,6 +48,8 @@ class ProgressRequest(BaseModel):
class InternalProgressResponse(BaseModel):
job: str = Field(default=None, title="Job name", description="Internal job name")
+ job_timestamp: str|None = Field(default=None, title="Job timestamp", description="Timestamp of the job start")
+ job_time: float|None = Field(default=None, title="Job start time", description="Time of the job start")
textinfo: str|None = Field(default=None, title="Info text", description="Info text used by WebUI.")
# status fields
active: bool = Field(title="Whether the task is being worked on right now")
@@ -122,7 +124,8 @@ def api_progress(req: ProgressRequest):
steps=steps,
batch_no=batch_no,
batch_count=batch_count,
- job_timestamp=shared.state.time_start,
+ job_timestamp=shared.state.job_timestamp,
+ job_time=shared.state.time_start,
eta=eta,
live_preview=live_preview,
id_live_preview=id_live_preview,
diff --git a/modules/rembg/lucida.py b/modules/rembg/lucida.py
new file mode 100644
index 000000000..9976cbc3a
--- /dev/null
+++ b/modules/rembg/lucida.py
@@ -0,0 +1,39 @@
+model = None
+repo_id = "egeorcun/lucida"
+
+
+def remove(image):
+ import torch
+ from PIL import Image
+ from torchvision import transforms
+ from transformers import AutoModelForImageSegmentation
+ from modules import devices
+
+ global model # pylint: disable=global-statement
+
+ if model is None:
+ model = AutoModelForImageSegmentation.from_pretrained(repo_id,
+ trust_remote_code=True,
+ dtype=torch.float32,
+ )
+ model.eval()
+
+ t = transforms.Compose([
+ transforms.Resize((1024, 1024)),
+ transforms.ToTensor(),
+ transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+ ])
+ model = model.to(device=devices.device)
+ with devices.inference_context():
+ input_tensor = t(image).unsqueeze(0).to(devices.device)
+ preds = model(input_tensor)[-1].sigmoid()
+ alpha = transforms.functional.resize(preds[0], image.size[::-1]).squeeze(0)
+ alpha = alpha.detach().cpu().numpy()
+ model = model.to(device=devices.cpu)
+
+ rgba = image.copy()
+ rgba.putalpha(Image.fromarray((255.0 * alpha).astype("uint8")))
+
+ if rgba is None:
+ return image
+ return rgba
diff --git a/modules/rembg/rembg_api.py b/modules/rembg/rembg_api.py
index e4ab757bd..0bb30f4dd 100644
--- a/modules/rembg/rembg_api.py
+++ b/modules/rembg/rembg_api.py
@@ -28,6 +28,9 @@ async def post_rembg(
if model == "ben2":
from modules.rembg import ben2
image = ben2.remove(input_image, refine=refine)
+ elif model == "lucida":
+ from modules.rembg import lucida
+ image = lucida.remove(input_image)
else:
dependencies()
import rembg
diff --git a/modules/scripts.py b/modules/scripts.py
index daca52161..3bb8d33a2 100644
--- a/modules/scripts.py
+++ b/modules/scripts.py
@@ -6,14 +6,16 @@ from modules.scripts_manager import * # pylint: disable=wildcard-import
scripts_txt2img = None
scripts_img2img = None
scripts_control = None
+scripts_video = None
scripts_current = None
scripts_postproc = None
def register_runners():
- global scripts_txt2img, scripts_img2img, scripts_control, scripts_current, scripts_postproc # pylint: disable=global-statement
+ global scripts_txt2img, scripts_img2img, scripts_control, scripts_video, scripts_current, scripts_postproc # pylint: disable=global-statement
scripts_txt2img = scripts_manager.scripts_txt2img
scripts_img2img = scripts_manager.scripts_img2img
scripts_control = scripts_manager.scripts_control
+ scripts_video = scripts_manager.scripts_video
scripts_current = scripts_manager.scripts_current
scripts_postproc = scripts_manager.scripts_postproc
diff --git a/modules/scripts_manager.py b/modules/scripts_manager.py
index d47d22197..96211a498 100644
--- a/modules/scripts_manager.py
+++ b/modules/scripts_manager.py
@@ -52,6 +52,8 @@ class Script:
alwayson = False
is_txt2img = False
is_img2img = False
+ is_control = False
+ is_video = False
api_info: ItemScript | None = None
group = None
infotext_fields: list | None = None
@@ -320,10 +322,11 @@ def load_scripts():
t.record(os.path.basename(scriptfile.basedir) if scriptfile.basedir != paths.script_path else scriptfile.filename)
sys.path = syspath
- global scripts_txt2img, scripts_img2img, scripts_control, scripts_postproc # pylint: disable=global-statement
+ global scripts_txt2img, scripts_img2img, scripts_control, scripts_video, scripts_postproc # pylint: disable=global-statement
scripts_txt2img = ScriptRunner('txt2img')
scripts_img2img = ScriptRunner('img2img')
scripts_control = ScriptRunner('control')
+ scripts_video = ScriptRunner('video')
scripts_postproc = scripts_postprocessing.ScriptPostprocessingRunner()
return t, time.time()-t0
@@ -373,12 +376,14 @@ class ScriptRunner:
self.inputs: list = [None]
self.time = 0
- def add_script(self, script_class, path, is_img2img, is_control):
+ def add_script(self, script_class, path, is_img2img, is_control, is_video):
try:
script = script_class()
script.filename = path
script.is_txt2img = not is_img2img
script.is_img2img = is_img2img
+ script.is_control = is_control
+ script.is_video = is_video
if path.startswith(paths.extensions_dir) and not path.startswith(paths.extensions_builtin_dir):
script.external = True
if is_control and script.external:
@@ -393,6 +398,8 @@ class ScriptRunner:
visibility = AlwaysVisible
else:
visibility = v1 or v2
+ elif is_video:
+ visibility = getattr(script, 'video_capable', False)
else:
visibility = script.show(script.is_img2img)
if visibility == AlwaysVisible:
@@ -406,7 +413,7 @@ class ScriptRunner:
log.error(f'Script initialize: {path} {e}')
errors.display(e, 'script')
- def initialize_scripts(self, is_img2img=False, is_control=False):
+ def initialize_scripts(self, is_img2img=False, is_control=False, is_video=False):
from modules import scripts_auto_postprocessing
self.scripts.clear()
@@ -428,14 +435,14 @@ class ScriptRunner:
except Exception:
sorted_scripts = scripts_data
for script_class, path, _basedir, _script_module in sorted_scripts:
- self.add_script(script_class, path, is_img2img, is_control)
+ self.add_script(script_class, path, is_img2img, is_control, is_video)
try:
sorted_scripts = sorted(self.auto_processing_scripts, key=lambda x: x.script_class().title().lower())
except Exception:
sorted_scripts = self.auto_processing_scripts
for script_class, path, _basedir, _script_module in sorted_scripts:
- self.add_script(script_class, path, is_img2img, is_control)
+ self.add_script(script_class, path, is_img2img, is_control, is_video)
def prepare_ui(self):
self.inputs = [None]
@@ -817,6 +824,7 @@ class ScriptRunner:
scripts_txt2img: ScriptRunner = None
scripts_img2img: ScriptRunner = None
scripts_control: ScriptRunner = None
+scripts_video: ScriptRunner = None
scripts_current: ScriptRunner = None
scripts_postproc: scripts_postprocessing.ScriptPostprocessingRunner = None
reload_scripts = load_scripts # compatibility alias
@@ -827,3 +835,4 @@ def reload_script_body_only():
scripts_txt2img.reload_sources(cache)
scripts_img2img.reload_sources(cache)
scripts_control.reload_sources(cache)
+ scripts_video.reload_sources(cache)
diff --git a/modules/sd_detect.py b/modules/sd_detect.py
index 96d5cc376..b6b1e72fb 100644
--- a/modules/sd_detect.py
+++ b/modules/sd_detect.py
@@ -175,6 +175,8 @@ def guess_by_name(fn, current_guess):
new_guess = 'JoyEdit'
elif 'sefi-image' in fn.lower():
new_guess = 'SeFi'
+ elif 'mage-flow' in fn.lower():
+ new_guess = 'MageFlow'
if debug_load:
log.trace(f'Autodetect: method=name file="{fn}" previous="{current_guess}" current="{new_guess}"')
return new_guess or current_guess
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 03d32b359..494c5bf1d 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -606,6 +606,10 @@ def load_diffuser_force(detected_model_type: str, checkpoint_info: CheckpointInf
from pipelines.model_sefi import load_sefi
sd_model = load_sefi(checkpoint_info, diffusers_load_config)
allow_post_quant = False
+ elif model_type in ['MageFlow']:
+ from pipelines.model_mageflow import load_mageflow
+ sd_model = load_mageflow(checkpoint_info, diffusers_load_config)
+ allow_post_quant = True
except Exception as e:
log.error(f'Load {op}: path="{checkpoint_info.path}" {e}')
errors.display(e, 'Load')
@@ -1393,9 +1397,9 @@ def set_diffuser_pipe(pipe, new_pipe_type):
add_noise_pred_to_diffusers_callback(new_pipe.pipe)
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
- log.debug(f"Pipeline class change: source={cls} target={new_pipe.__class__.__name__} device={pipe.device} fn={fn}") # pylint: disable=protected-access
+ log.debug(f"Pipeline class change: source={cls} target={new_pipe.__class__.__name__} fn={fn}") # pylint: disable=protected-access
- if shared.opts.diffusers_offload_mode == 'none':
+ if shared.opts.diffusers_offload_mode == 'none' and hasattr(pipe, 'device'):
move_model(new_pipe, pipe.device)
else:
set_diffuser_offload(new_pipe, op='model')
@@ -1428,7 +1432,7 @@ def add_noise_pred_to_diffusers_callback(pipe):
def get_native(pipe: diffusers.DiffusionPipeline):
- if hasattr(pipe, "vae") and hasattr(pipe.vae.config, "sample_size"):
+ if hasattr(pipe, "vae") and hasattr(pipe.vae, "config") and hasattr(pipe.vae.config, "sample_size"):
size = pipe.vae.config.sample_size # Stable Diffusion
elif hasattr(pipe, "movq") and hasattr(pipe.movq.config, "sample_size"):
size = pipe.movq.config.sample_size # Kandinsky
@@ -1548,14 +1552,17 @@ def clear_caches(full: bool = False):
lora_common.loaded_networks.clear()
lora_common.previously_loaded_networks.clear()
lora_load.lora_cache.clear()
+ fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
+ log.debug(f'Cache clear: full={full} fn={fn}')
if full:
- log.debug('Cache clear')
sd_offload.offload_hook_instance = None
def unload_model_weights(op='model'):
fn = f'{sys._getframe(2).f_code.co_name}:{sys._getframe(1).f_code.co_name}' # pylint: disable=protected-access
- clear_caches(full=True)
+ if model_data.sd_model or model_data.sd_refiner:
+ clear_caches(full=True)
+ devices.torch_reset()
if shared.compiled_model_state is not None:
shared.compiled_model_state.compiled_cache.clear()
shared.compiled_model_state.req_cache.clear()
diff --git a/modules/sd_models_utils.py b/modules/sd_models_utils.py
index 7ff223638..fa4712ddf 100644
--- a/modules/sd_models_utils.py
+++ b/modules/sd_models_utils.py
@@ -72,7 +72,7 @@ def path_to_repo(checkpoint_info: CheckpointInfo | str):
for opt in [shared.opts.ckpt_dir, shared.opts.diffusers_dir, shared.opts.hfcache_dir]:
remove_prefix.append(opt.replace('\\', '/'))
try:
- relative = os.path.relpath(opt, start=shared.opts.models_dir).replace('\\', '/')
+ relative = os.path.relpath(opt, start=shared.models_path).replace('\\', '/')
if not relative.startswith('.'):
remove_prefix.append(relative)
except Exception:
diff --git a/modules/sd_offload.py b/modules/sd_offload.py
index c6d2cdc9a..246ac6e2f 100644
--- a/modules/sd_offload.py
+++ b/modules/sd_offload.py
@@ -250,12 +250,12 @@ class OffloadHook(accelerate.hooks.ModelHook):
for module_name in get_module_names(pipe):
module_instance = getattr(pipe, module_name, None)
module_cls = module_instance.__class__.__name__
- if (module_instance is not None) and (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(module_instance.device, devices.cpu)):
+ if (module_instance is not None) and (_id != id(module_instance)) and (module_cls not in self.offload_never) and (not devices.same_device(getattr(module_instance, "device", devices.cpu), devices.cpu)):
apply_balanced_offload_to_module(module_instance, op='pre')
self.last_cls = module.__class__.__name__
process_timer.add('offload', time.time() - t0)
- if not devices.same_device(module.device, devices.device): # move-to-device
+ if not devices.same_device(getattr(module, "device", devices.cpu), devices.device): # move-to-device
t0 = time.time()
device_index = torch.device(devices.device).index
if device_index is None:
@@ -301,7 +301,7 @@ class OffloadHook(accelerate.hooks.ModelHook):
for _i, pipe in enumerate(get_pipe_variants()):
for module_name in get_module_names(pipe):
module_instance = getattr(pipe, module_name, None)
- log.trace(f'Offload: type=balanced op=pre:status forward={module.__class__.__name__} module={module_name} class={module_instance.__class__.__name__} pipe={_i} device={module_instance.device} dtype={module_instance.dtype}')
+ log.trace(f'Offload: type=balanced op=pre:status forward={module.__class__.__name__} module={module_name} class={module_instance.__class__.__name__} pipe={_i} device={getattr(module_instance, "device", devices.cpu)} dtype={module_instance.dtype}')
self.last_pre = _id
return args, kwargs
diff --git a/modules/sd_offload_aux.py b/modules/sd_offload_aux.py
index 26c8c9fa7..e8cf54d50 100644
--- a/modules/sd_offload_aux.py
+++ b/modules/sd_offload_aux.py
@@ -22,6 +22,8 @@ aux_models: dict[str, AuxModel] = {}
def register_aux(name: str, model: torch.nn.Module) -> None:
+ if model is None:
+ return
size = sum(p.numel() * p.element_size() for p in model.parameters()) / 1024**3
aux_models[name] = AuxModel(model=model, name=name, size=size)
debug_move(f'Offload: type=aux op=register name={name} size={size:.3f}')
@@ -42,6 +44,8 @@ def evict_aux(exclude: str | None = None, reason: str = 'evict') -> None:
def _do_move_to_cpu(model, op_label, size):
+ if model is None:
+ return
if shared.opts.diffusers_offload_streams:
global move_stream # pylint: disable=global-statement
if move_stream is None:
diff --git a/modules/sd_samplers_common.py b/modules/sd_samplers_common.py
index d964a4bf8..2a9e8c863 100644
--- a/modules/sd_samplers_common.py
+++ b/modules/sd_samplers_common.py
@@ -1,3 +1,4 @@
+import os
import time
import threading
from collections import namedtuple
@@ -14,6 +15,7 @@ approximation_indexes = { "Simple": 0, "Approximate": 1, "TAESD": 2, "Full VAE":
flow_models = ['f1', 'f2', 'sd3', 'lumina', 'auraflow', 'sana', 'zimage', 'lumina2', 'cogview4', 'h1', 'cosmos', 'anima', 'chroma', 'omnigen', 'omnigen2', 'longcat', 'ideogram4', 'krea2']
warned = False
queue_lock = threading.Lock()
+debug = os.environ.get('SD_PREVIEW_DEBUG', None) is not None
def warn_once(message):
@@ -35,10 +37,12 @@ def setup_img2img_steps(p, steps=None):
return steps, t_enc
-def single_sample_to_image(sample, approximation=None):
+def single_sample_to_image(sample, approximation=None, fast=False):
with queue_lock:
t0 = time.time()
approximation = approximation or shared.opts.show_progress_type
+ if debug:
+ log.debug(f'Preview sample: shape={list(sample.shape)} dtype={sample.dtype} method={approximation}')
try:
if (sample.dtype == torch.bfloat16) and (approximation in ["Simple", "Approximate"]):
sample = sample.to(torch.float16)
@@ -59,7 +63,7 @@ def single_sample_to_image(sample, approximation=None):
sample = torch.nn.functional.interpolate(sample.unsqueeze(0), scale_factor=[scale, scale], mode='bilinear', align_corners=False)[0]
except Exception:
pass
- x_sample = sd_vae_taesd.decode(sample)
+ x_sample = sd_vae_taesd.decode(sample, fast=fast)
# x_sample = (1.0 + x_sample) / 2.0 # preview requires smaller range
elif shared.sd_model_type == 'sc' and approximation != "Full":
x_sample = sd_vae_stablecascade.decode(sample)
@@ -98,8 +102,8 @@ def sample_to_image(samples, index=0, approximation=None):
return single_sample_to_image(samples[index], approximation)
-def samples_to_image_grid(samples, approximation=None):
- return images.image_grid([single_sample_to_image(sample, approximation) for sample in samples])
+def samples_to_image_grid(samples, approximation=None, fast=False):
+ return images.image_grid([single_sample_to_image(sample, approximation, fast=fast) for sample in samples])
def store_latent(decoded):
diff --git a/modules/sdnq/common.py b/modules/sdnq/common.py
index 30d8a82e7..94db1cbe9 100644
--- a/modules/sdnq/common.py
+++ b/modules/sdnq/common.py
@@ -6,7 +6,7 @@ import torch
from modules import shared
-sdnq_version = "0.2.3"
+sdnq_version = "0.2.4"
sdnq_keys = {"weight", "scale", "zero_point", "svd_up", "svd_down"}
torch_version = torch.__version__[:4]
@@ -357,9 +357,8 @@ if use_torch_compile:
kwargs["fullgraph"] = True
if kwargs.get("dynamic", None) is None:
kwargs["dynamic"] = False
- if torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] >= 12):
- if kwargs.get("recompile_limit", None) is None:
- kwargs["recompile_limit"] = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
+ if (torch_version[0] > 2 or (torch_version[0] == 2 and torch_version[1] >= 12)) and kwargs.get("recompile_limit", None) is None:
+ kwargs["recompile_limit"] = max(8192, getattr(torch._dynamo.config, "recompile_limit", 0))
if os.environ.get("SDNQ_COMPILE_KWARGS", None) is not None:
for key, value in json.loads(os.environ.get("SDNQ_COMPILE_KWARGS")).items():
kwargs[key] = value
diff --git a/modules/sdnq/file_loader.py b/modules/sdnq/file_loader.py
index d5fe53574..7ea8e1fbe 100644
--- a/modules/sdnq/file_loader.py
+++ b/modules/sdnq/file_loader.py
@@ -19,7 +19,7 @@ def load_safetensors(files: list[str], state_dict: dict | None = None, key_mappi
state_dict = {}
for fn in files:
with safe_open(fn, framework="pt", device=str(device)) as f:
- for key in f.keys():
+ for key in f:
state_dict[map_keys(key, key_mapping)] = f.get_tensor(key)
diff --git a/modules/sdnq/kernel_wrappers.py b/modules/sdnq/kernel_wrappers.py
index a27d2d994..9a9751fbc 100644
--- a/modules/sdnq/kernel_wrappers.py
+++ b/modules/sdnq/kernel_wrappers.py
@@ -25,7 +25,7 @@ if os.environ.get("SDNQ_ALLOW_FP8_COMPILE", None) is None:
else:
is_fp8_compile_supported = True
else:
- is_fp8_compile_supported = os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"}
+ is_fp8_compile_supported = bool(os.environ.get("SDNQ_ALLOW_FP8_COMPILE", "0").lower() not in {"0", "false", "no"})
if devices.backend == "rocm":
gfx_version = devices.get_hip_agent().gfx_version
@@ -38,82 +38,56 @@ if devices.backend in {"ipex", "xpu"}:
else:
is_alchemist_or_igpu = False
-if os.environ.get("SDNQ_USE_OPENVINO_MM", None) is None:
- use_openvino_mm = bool(devices.backend in {"cpu", "openvino"})
-else:
- use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "0").lower() not in {"0", "false", "no"})
-
if os.environ.get("SDNQ_USE_TRITON_MM", None) is None:
use_triton_mm = bool(not is_alchemist_or_igpu and (devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"}))
else:
use_triton_mm = bool(os.environ.get("SDNQ_USE_TRITON_MM", "0").lower() not in {"0", "false", "no"})
-use_triton_scaled_mm = bool(use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").lower() not in {"0", "false", "no"})
if os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", None) is None:
# row-wise FP8 only exist on H100 hardware, sdnq will use software row-wise with tensorwise hardware with this setting
use_tensorwise_fp8_matmul = bool(devices.backend != "cuda" or (devices.backend == "cuda" and torch.cuda.get_device_capability(devices.device) < (9,0)))
else:
- use_tensorwise_fp8_matmul = os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"}
+ use_tensorwise_fp8_matmul = bool(os.environ.get("SDNQ_USE_TENSORWISE_FP8_MM", "0").lower() not in {"0", "false", "no"})
-if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
- use_contiguous_int8_mm = bool(use_openvino_mm or is_rdna2_and_older or devices.backend in {"ipex", "xpu", "mps", "openvino", "zluda"})
- use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
- use_contiguous_fp8_mm = use_contiguous_fp16_mm
-else:
- use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
- use_contiguous_fp16_mm = use_contiguous_int8_mm
- use_contiguous_fp8_mm = use_contiguous_fp16_mm
-
-
-int_mm_func = None
-fp_mm_func = None
-fp8_mm_func = None
-int_scaled_mm_func = None
-fp_scaled_mm_func = None
-fp8_scaled_mm_func = None
+use_openvino_mm = bool(os.environ.get("SDNQ_USE_OPENVINO_MM", "1").lower() not in {"0", "false", "no"})
+use_triton_scaled_mm = bool(use_triton_mm and os.environ.get("SDNQ_USE_TRITON_SCALED_MM", "1").lower() not in {"0", "false", "no"})
if use_openvino_mm:
try:
from .kernels.openvino_mm import openvino_int_mm, openvino_fp_mm
- int_mm_func = openvino_int_mm
- fp_mm_func = openvino_fp_mm
- except Exception:
+ except Exception as e:
use_openvino_mm = False
-elif use_triton_mm:
+ openvino_int_mm = None
+ openvino_fp_mm = None
+ shared.log.warning(f"SDNQ: OpenVINO MM kernels are not available! Falling back to PyTorch Eager kernels for CPU device. Error message: {e}")
+else:
+ openvino_int_mm = None
+ openvino_fp_mm = None
+
+
+if use_triton_mm:
try:
from .kernels.triton_mm import sdnq_triton_mm
- int_mm_func = sdnq_triton_mm
- fp_mm_func = sdnq_triton_mm
if is_fp8_mm_supported:
- fp8_mm_func = sdnq_triton_mm
use_tensorwise_fp8_matmul = True
- if use_triton_scaled_mm:
- from .kernels.triton_scaled_mm import sdnq_scaled_mm
- int_scaled_mm_func = sdnq_scaled_mm
- fp_scaled_mm_func = sdnq_scaled_mm
- if is_fp8_mm_supported:
- fp8_scaled_mm_func = sdnq_scaled_mm
except Exception as e:
use_triton_mm = False
- use_triton_scaled_mm = False
- shared.log.warning(f"SDNQ: Triton kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
+ sdnq_triton_mm = None
+ shared.log.warning(f"SDNQ: Triton MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
+else:
+ sdnq_triton_mm = None
-if (
- fp_mm_func is None and not is_alchemist_or_igpu
- and devices.backend in {"cuda", "rocm", "ipex", "xpu", "zluda"}
- and os.environ.get("SDNQ_USE_TRITON_MM", "1").lower() not in {"0", "false", "no"}
-):
+if use_triton_scaled_mm:
try:
- from .kernels.triton_mm import sdnq_triton_mm
- fp_mm_func = sdnq_triton_mm
- if use_triton_scaled_mm:
- from .kernels.triton_scaled_mm import sdnq_scaled_mm
- fp_scaled_mm_func = sdnq_scaled_mm
- except Exception:
- use_triton_mm = False
+ from .kernels.triton_scaled_mm import sdnq_scaled_mm
+ except Exception as e:
use_triton_scaled_mm = False
+ sdnq_scaled_mm = None
+ shared.log.warning(f"SDNQ: Triton Scaled MM kernels are not available! Falling back to PyTorch Eager kernels. Error message: {e}")
+else:
+ sdnq_scaled_mm = None
if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None:
@@ -121,9 +95,24 @@ if os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", None) is None:
else:
include_mm_kernel_in_compile = bool(os.environ.get("SDNQ_INCLUDE_MM_KERNEL_IN_COMPILE", "0").lower() not in {"0", "false", "no"})
+if os.environ.get("SDNQ_USE_CONTIGUOUS_MM", None) is None:
+ use_contiguous_int8_mm = bool(is_rdna2_and_older or devices.backend in {"ipex", "xpu", "cpu", "mps", "openvino", "zluda"})
+ use_contiguous_fp16_mm = bool(use_contiguous_int8_mm or devices.backend == "rocm")
+ use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported and use_triton_mm)
+else:
+ use_contiguous_int8_mm = bool(os.environ.get("SDNQ_USE_CONTIGUOUS_MM", "0").lower() not in {"0", "false", "no"})
+ use_contiguous_fp16_mm = use_contiguous_int8_mm
+ use_contiguous_fp8_mm = use_contiguous_fp16_mm and (is_fp8_mm_supported and use_triton_mm)
+
+
+def int_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
+ return torch._int_mm(a,b).to(dtype=out_dtype)
+
+
+def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ dummy_input_scale = torch.ones(1, device=a.device, dtype=torch.float32)
+ return torch._scaled_mm(a, b, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=out_dtype)
-def fp_mm_torch_cuda(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- return torch.mm(a,b, out_dtype=out_dtype)
def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if b.dtype == torch.float8_e4m3fn:
@@ -136,59 +125,91 @@ def fp_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch
return torch.mm(a,b).to(dtype=torch.float32).mul_(fp16_scale).to(dtype=out_dtype)
-if int_mm_func is None:
- def int_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
- return torch._int_mm(a,b).to(dtype=out_dtype)
- int_mm_func = int_mm_torch
-
-
-if fp_mm_func is None:
- if devices.backend == "cuda":
- fp_mm_func = fp_mm_torch_cuda
+def int_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if bias is None:
+ return int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
else:
- fp_mm_func = fp_mm_torch
+ return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
-if fp8_mm_func is None:
+if use_tensorwise_fp8_matmul or not is_fp8_mm_supported:
+ def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if bias is None:
+ return fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+ else:
+ return torch.addcmul(bias, fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+else:
+ def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if bias is not None and bias.ndim != 1:
+ return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=None, out_dtype=out_dtype).add_(bias)
+ else:
+ return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=bias.to(dtype=out_dtype) if bias is not None else None, out_dtype=out_dtype)
+
+
+def fp_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if bias is None:
+ return fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+ else:
+ return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+
+
+def int_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.int32) -> torch.FloatTensor:
+ if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_triton_mm(a, b, out_dtype=out_dtype)
+ elif openvino_int_mm is not None and a.device.type == "cpu":
+ return openvino_int_mm(a, b, out_dtype=out_dtype)
+ else:
+ return int_mm_torch(a, b, out_dtype=out_dtype)
+
+
+def fp8_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
if is_fp8_mm_supported:
- def fp8_mm_torch(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- dummy_input_scale = torch.ones(1, device=a.device, dtype=torch.float32)
- return torch._scaled_mm(a, b, scale_a=dummy_input_scale, scale_b=dummy_input_scale, bias=None, out_dtype=out_dtype)
- fp8_mm_func = fp8_mm_torch
- use_contiguous_fp8_mm = False
- else:
- fp8_mm_func = fp_mm_torch
-
-
-if int_scaled_mm_func is None:
- def int_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- if bias is None:
- return int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
+ if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_triton_mm(a, b, out_dtype=out_dtype)
+ elif openvino_fp_mm is not None and a.device.type == "cpu":
+ return openvino_fp_mm(a, b, out_dtype=out_dtype)
else:
- return torch.addcmul(bias, int_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
- int_scaled_mm_func = compile_func(int_scaled_mm_torch)
-
-
-if fp_scaled_mm_func is None:
- def fp_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- if bias is None:
- return fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
- else:
- return torch.addcmul(bias, fp_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
- fp_scaled_mm_func = compile_func(fp_scaled_mm_torch)
-
-
-if fp8_scaled_mm_func is None:
- if use_tensorwise_fp8_matmul or not is_fp8_mm_supported:
- def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- if bias is None:
- return fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a).mul_(scale_b).to(dtype=out_dtype)
- else:
- return torch.addcmul(bias, fp8_mm_func(a,b, out_dtype=scale_a.dtype).mul_(scale_a), scale_b).to(dtype=out_dtype)
+ return fp8_mm_torch(a, b, out_dtype=out_dtype)
else:
- def fp8_scaled_mm_torch(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
- if bias is not None and bias.ndim != 1:
- return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=None, out_dtype=out_dtype).add_(bias)
- else:
- return torch._scaled_mm(a, b, scale_a=scale_a, scale_b=scale_b, bias=bias.to(dtype=out_dtype) if bias is not None else None, out_dtype=out_dtype)
- fp8_scaled_mm_func = compile_func(fp8_scaled_mm_torch)
+ if openvino_fp_mm is not None and a.device.type == "cpu":
+ return openvino_fp_mm(a, b, out_dtype=out_dtype)
+ else:
+ return fp_mm_torch(a, b, out_dtype=out_dtype)
+
+
+def fp_mm_func(a: torch.Tensor, b: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if sdnq_triton_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_triton_mm(a, b, out_dtype=out_dtype)
+ elif openvino_fp_mm is not None and a.device.type == "cpu":
+ return openvino_fp_mm(a, b, out_dtype=out_dtype)
+ else:
+ return fp_mm_torch(a, b, out_dtype=out_dtype)
+
+
+def int_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+ else:
+ return int_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+
+
+def fp8_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if is_fp8_mm_supported and sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+ else:
+ return fp8_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+
+
+def fp_scaled_mm_func(a: torch.Tensor, b: torch.Tensor, scale_a: torch.Tensor, scale_b: torch.Tensor, bias: torch.FloatTensor | None = None, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+ if sdnq_scaled_mm is not None and a.device.type in {"cuda", "xpu"}:
+ return sdnq_scaled_mm(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+ else:
+ return fp_scaled_mm_torch(a, b, scale_a, scale_b, bias=bias, out_dtype=out_dtype)
+
+
+int_mm_torch = compile_func(int_mm_torch)
+fp8_mm_torch = compile_func(fp8_mm_torch)
+fp_mm_torch = compile_func(fp_mm_torch)
+int_scaled_mm_torch = compile_func(int_scaled_mm_torch)
+fp8_scaled_mm_torch = compile_func(fp8_scaled_mm_torch)
+fp_scaled_mm_torch = compile_func(fp_scaled_mm_torch)
diff --git a/modules/sdnq/kernels/openvino_mm.py b/modules/sdnq/kernels/openvino_mm.py
index 3e850810c..b0832b009 100644
--- a/modules/sdnq/kernels/openvino_mm.py
+++ b/modules/sdnq/kernels/openvino_mm.py
@@ -2,23 +2,32 @@ import os
import torch
import openvino as ov
from openvino import opset16 as ov_ops
-from openvino.properties import hint as ov_hints
-core = ov.Core()
-
-OV_DEVICE: str = os.environ.get("SDNQ_OPENVINO_DEVICE", "HETERO:NPU,CPU" if "NPU" in core.get_available_devices() else "CPU")
+OV_CORE = None
+OV_DEVICE: str = None
OV_COMPILED_CACHE: dict[tuple[str, tuple[int,int] | None, str, tuple[int,int] | None], tuple[ov.InferRequest, str]] = {}
-if OV_DEVICE == "NPU":
- OV_DEVICE = "HETERO:NPU,CPU"
-for ov_device in core.get_available_devices():
- core.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
+
+def get_ov_core():
+ global OV_CORE, OV_DEVICE # pylint: disable=global-statement
+ if OV_CORE is not None:
+ return OV_CORE, OV_DEVICE
+ from openvino.properties import hint as ov_hints
+ OV_CORE = ov.Core()
+ OV_DEVICE = os.environ.get("SDNQ_OPENVINO_DEVICE", "HETERO:NPU,CPU" if "NPU" in OV_CORE.get_available_devices() else "CPU")
+ if OV_DEVICE == "NPU":
+ OV_DEVICE = "HETERO:NPU,CPU"
+ for ov_device in OV_DEVICE.removeprefix("HETERO:").split(","):
+ if ov_device != "NPU":
+ OV_CORE.set_property(ov_device, {ov_hints.execution_mode: ov_hints.ExecutionMode.ACCURACY})
+ return OV_CORE, OV_DEVICE
-def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_name: str, out_dtype: torch.dtype = torch.float32) -> torch.FloatTensor:
+def ov_mm(infer_request: ov.InferRequest, out_name: str, A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
C = torch.empty((A.shape[0], B.shape[-1]), device="cpu", dtype=torch.float32)
- infer_request.set_tensor("A", ov.Tensor(A.detach().contiguous().to("cpu").numpy(), shared_memory=True))
- infer_request.set_tensor("B", ov.Tensor(B.detach().contiguous().to("cpu").numpy(), shared_memory=True))
+ A, B = A.contiguous(), B.contiguous()
+ infer_request.set_tensor("A", ov.Tensor(A.detach().to("cpu").numpy(), shared_memory=True))
+ infer_request.set_tensor("B", ov.Tensor(B.detach().to("cpu").numpy(), shared_memory=True))
infer_request.set_tensor(out_name, ov.Tensor(C.numpy(), shared_memory=True))
infer_request.infer()
C = C.to(A.device, dtype=out_dtype)
@@ -27,20 +36,22 @@ def ov_mm(A: torch.Tensor, B: torch.Tensor, infer_request: ov.InferRequest, out_
@torch.library.custom_op("sdnq::openvino_int_mm", mutates_args=())
def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
- if "GPU" not in OV_DEVICE:
- cache_key = (OV_DEVICE, "int8", Tensor_A.shape, Tensor_B.shape)
+ ov_core, ov_device = get_ov_core()
+ if "GPU" not in ov_device:
+ cache_key = (ov_device, "int8", Tensor_A.shape, Tensor_B.shape)
else:
- cache_key = (OV_DEVICE, "int8", None, None)
+ cache_key = (ov_device, "int8", None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
- return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
+ return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
- if "GPU" not in OV_DEVICE:
+ if "GPU" not in ov_device:
shape_a = ov.Shape(Tensor_A.shape)
shape_b = ov.Shape(Tensor_B.shape)
else:
shape_a = ov.PartialShape([-1,-1])
shape_b = ov.PartialShape([-1,-1])
+
input_a = ov_ops.parameter(shape_a, ov.Type.i8, name="A")
input_b = ov_ops.parameter(shape_b, ov.Type.i8, name="B")
a = ov_ops.convert(input_a, ov.Type.f32)
@@ -52,7 +63,7 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: t
b = ov_ops.fake_quantize(b, low, high, low, high, 256)
# NPU uses FP16 x INT8 -> FP16 instead of INT8 x INT8 -> INT32 and FP16 output overflows
- if "NPU" in OV_DEVICE:
+ if "NPU" in ov_device:
fp16_scale = 0.25012213 * Tensor_B.shape[-2]
in_scale = ov_ops.constant(fp16_scale ** 0.5, dtype=ov.Type.f32)
out_scale = ov_ops.constant(fp16_scale, dtype=ov.Type.f32, name="out_scale_const")
@@ -64,18 +75,18 @@ def openvino_int_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: t
out = ov_ops.matmul(a, b, False, False)
ov_model = ov.Model([out], [input_a, input_b], "ov_int8_mm")
- if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications
+ if "NPU" in ov_device: # NPU can't use FP32 for regular multiplications
for node in ov_model.get_ops():
if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
node.get_rt_info()["affinity"] = "CPU"
else:
node.get_rt_info()["affinity"] = "NPU"
- ov_model = core.compile_model(ov_model, OV_DEVICE)
+ ov_model = ov_core.compile_model(ov_model, ov_device)
infer_request = ov_model.create_infer_request()
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
- return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
+ return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
@openvino_int_mm.register_fake
def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
@@ -84,24 +95,26 @@ def openvino_int_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtyp
@torch.library.custom_op("sdnq::openvino_fp_mm", mutates_args=())
def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
+ ov_core, ov_device = get_ov_core()
mm_dtype = "fp16" if Tensor_B.dtype == torch.float16 else "fp8"
if mm_dtype == "fp8":
Tensor_A = Tensor_A.to(dtype=torch.float16)
Tensor_B = Tensor_B.to(dtype=torch.float16)
- if "GPU" not in OV_DEVICE:
- cache_key = (OV_DEVICE, mm_dtype, Tensor_A.shape, Tensor_B.shape)
+ if "GPU" not in ov_device:
+ cache_key = (ov_device, mm_dtype, Tensor_A.shape, Tensor_B.shape)
else:
- cache_key = (OV_DEVICE, mm_dtype, None, None)
+ cache_key = (ov_device, mm_dtype, None, None)
infer_request, out_name = OV_COMPILED_CACHE.get(cache_key, (None, None))
if infer_request is not None:
- return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
+ return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
- if "GPU" not in OV_DEVICE:
+ if "GPU" not in ov_device:
shape_a = ov.Shape(Tensor_A.shape)
shape_b = ov.Shape(Tensor_B.shape)
else:
shape_a = ov.PartialShape([-1,-1])
shape_b = ov.PartialShape([-1,-1])
+
input_a = ov_ops.parameter(shape_a, ov.Type.f16, name="A")
input_b = ov_ops.parameter(shape_b, ov.Type.f16, name="B")
a = ov_ops.convert(input_a, ov.Type.f32)
@@ -124,18 +137,18 @@ def openvino_fp_mm(Tensor_A: torch.Tensor, Tensor_B: torch.Tensor, out_dtype: to
out = ov_ops.multiply(ov_ops.convert(out, ov.Type.f32), out_scale, name="out_scale")
ov_model = ov.Model([out], [input_a, input_b], "ov_fp_mm")
- if "NPU" in OV_DEVICE: # NPU can't use FP32 for regular multiplications
+ if "NPU" in ov_device: # NPU can't use FP32 for regular multiplications
for node in ov_model.get_ops():
if node.get_friendly_name() in {"out_scale", "out_scale_const"}:
node.get_rt_info()["affinity"] = "CPU"
else:
node.get_rt_info()["affinity"] = "NPU"
- ov_model = core.compile_model(ov_model, OV_DEVICE)
+ ov_model = ov_core.compile_model(ov_model, ov_device)
infer_request = ov_model.create_infer_request()
out_name = ov_model.outputs[0]
OV_COMPILED_CACHE[cache_key] = (infer_request, out_name)
- return ov_mm(Tensor_A, Tensor_B, infer_request, out_name, out_dtype=out_dtype)
+ return ov_mm(infer_request, out_name, Tensor_A, Tensor_B, out_dtype=out_dtype)
@openvino_fp_mm.register_fake
def openvino_fp_mm_fake(A: torch.Tensor, B: torch.Tensor, out_dtype: torch.dtype = torch.float32) -> torch.Tensor:
diff --git a/modules/sdnq/layers/__init__.py b/modules/sdnq/layers/__init__.py
index 6835d059e..2e66f7718 100644
--- a/modules/sdnq/layers/__init__.py
+++ b/modules/sdnq/layers/__init__.py
@@ -30,7 +30,7 @@ class SDNQLayer(torch.nn.Module):
return self.forward_func(self, *args, **kwargs)
def __repr__(self) -> str:
- return f"{self.__class__.__name__}(original_class={self.original_class} forward_func={self.forward_func} sdnq_dequantizer={repr(getattr(self, 'sdnq_dequantizer', None))})"
+ return f"{self.__class__.__name__}(original_class={self.original_class} forward_func={self.forward_func} sdnq_dequantizer={getattr(self, 'sdnq_dequantizer', None)})"
class SDNQLinear(SDNQLayer, torch.nn.Linear):
diff --git a/modules/sdnq/loader.py b/modules/sdnq/loader.py
index 276c609cd..60235af69 100644
--- a/modules/sdnq/loader.py
+++ b/modules/sdnq/loader.py
@@ -168,9 +168,8 @@ def load_sdnq_model(
# older transformers case, handle known models manually
if model.__class__.__name__ in {"T5EncoderModel", "UMT5EncoderModel"} and "encoder.embed_tokens.weight" not in state_dict:
state_dict["encoder.embed_tokens.weight"] = state_dict["shared.weight"]
- elif model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict:
- if "model.embed_tokens.weight" in state_dict:
- state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"]
+ elif model.__class__.__name__ in {"Qwen3ForCausalLM"} and "lm_head.weight" not in state_dict and "model.embed_tokens.weight" in state_dict:
+ state_dict["lm_head.weight"] = state_dict["model.embed_tokens.weight"]
model.load_state_dict(state_dict, assign=True)
del state_dict
diff --git a/modules/sdnq/quant_utils.py b/modules/sdnq/quant_utils.py
index 4e52295bd..a711cfe64 100644
--- a/modules/sdnq/quant_utils.py
+++ b/modules/sdnq/quant_utils.py
@@ -130,6 +130,8 @@ def rotate_hadamard(weight: torch.Tensor, group_size: int = 256, hadamard: torch
hadamard = get_hadamard(group_size, dtype=weight.dtype, device=weight.device)
else:
group_size = hadamard.shape[-1]
+ if hadamard.dtype != weight.dtype:
+ hadamard = hadamard.to(dtype=weight.dtype)
if is_conv:
weight_shape = list(weight.shape)[1:]
weight = weight.flatten(1,-1)
diff --git a/modules/sdnq/quantizer.py b/modules/sdnq/quantizer.py
index 66a4131fc..e93e8b76e 100644
--- a/modules/sdnq/quantizer.py
+++ b/modules/sdnq/quantizer.py
@@ -3,6 +3,7 @@
from dataclasses import dataclass
from enum import Enum
+import os
import torch
from transformers.quantizers import HfQuantizer
@@ -53,6 +54,10 @@ from .layers import get_sdnq_wrapper_class
from .common import sdnq_version as current_sdnq_version
+from diffusers import __version__ as diffusers_version_str # pylint: disable=ungrouped-imports,wrong-import-order
+diffusers_version = [int(i) for i in diffusers_version_str.split(".")[:3]]
+
+
class QuantizationMethod(str, Enum):
SDNQ = "sdnq"
SDNQ_TRAINING = "sdnq_training"
@@ -276,7 +281,7 @@ def sdnq_quantize_layer_weight_dynamic(
param_name: str | None = None,
torch_dtype: torch.dtype | None = None,
quantization_config: "SDNQConfig" = None,
-) -> None | tuple[SDNQDequantizer, dict[str, torch.Tensor]]:
+) -> tuple[SDNQDequantizer, dict[str, torch.Tensor]] | None:
if torch_dtype is None:
torch_dtype = weight.dtype
if dynamic_loss_threshold is None or dynamic_loss_threshold < 0:
@@ -375,7 +380,7 @@ def sdnq_quantize_layer_weight_dynamic(
if quantization_loss <= dynamic_loss_threshold:
del original_weight_fp32
if quantization_config is not None:
- if sdnq_dequantizer.weights_dtype not in quantization_config.modules_dtype_dict.keys():
+ if sdnq_dequantizer.weights_dtype not in quantization_config.modules_dtype_dict:
quantization_config.modules_dtype_dict[sdnq_dequantizer.weights_dtype] = [param_name]
else:
quantization_config.modules_dtype_dict[sdnq_dequantizer.weights_dtype].append(param_name)
@@ -491,6 +496,7 @@ def sdnq_post_load_quant(
non_blocking: bool = False,
add_skip_keys:bool = True,
minimum_allowed_numel: int = 16384,
+ minimum_allowed_channel_size: int = 32,
modules_to_not_convert: list[str] | None = None,
modules_to_not_use_matmul: list[str] | None = None,
modules_dtype_dict: dict[str, list[str]] | None = None,
@@ -530,6 +536,7 @@ def sdnq_post_load_quant(
non_blocking=non_blocking,
add_skip_keys=add_skip_keys,
minimum_allowed_numel=minimum_allowed_numel,
+ minimum_allowed_channel_size=minimum_allowed_channel_size,
modules_to_not_convert=modules_to_not_convert,
modules_to_not_use_matmul=modules_to_not_use_matmul,
modules_dtype_dict=modules_dtype_dict,
@@ -570,7 +577,7 @@ class SDNQQuantize:
missing_keys: list[str] | None = None,
**kwargs,
) -> dict[str, torch.Tensor]:
- _module_name, value = tuple(input_dict.items())[0]
+ _module_name, value = next(iter(input_dict.items()))
value = value[0]
self.hf_quantizer.create_quantized_param(model, value, full_layer_name, value.device)
param, name = get_module_from_name(model, full_layer_name)
@@ -868,6 +875,8 @@ class SDNQConfig(QuantizationConfigMixin):
Disabling this option won't add model specific keys to modules_to_not_convert, modules_to_not_use_matmul and modules_dtype_dict.
minimum_allowed_numel (`int`, *optional*, defaults to `16384`):
Layers that have less than `minimum_allowed_numel` elements in them will be skipped and added to `modules_to_not_convert`.
+ minimum_allowed_channel_size (`int`, *optional*, defaults to `32`):
+ Layers that have less than `minimum_allowed_channel_size` channels in them will be skipped and added to `modules_to_not_convert`.
modules_to_not_convert (`list`, *optional*, default to `None`):
The list of modules to not quantize. Useful for quantizing models that explicitly require to have some
modules left in their original precision (e.g. Whisper encoder, Llava encoder, Mixtral gate layers).
@@ -912,6 +921,7 @@ class SDNQConfig(QuantizationConfigMixin):
non_blocking: bool = False,
add_skip_keys: bool = True,
minimum_allowed_numel: int = 16384,
+ minimum_allowed_channel_size: int = 32,
modules_to_not_convert: list[str] | None = None,
modules_to_not_use_matmul: list[str] | None = None,
modules_dtype_dict: dict[str, list[str]] | None = None,
@@ -944,6 +954,7 @@ class SDNQConfig(QuantizationConfigMixin):
self.non_blocking = non_blocking
self.add_skip_keys = add_skip_keys
self.minimum_allowed_numel = minimum_allowed_numel
+ self.minimum_allowed_channel_size = minimum_allowed_channel_size
self.modules_to_not_convert = modules_to_not_convert
self.modules_to_not_use_matmul = modules_to_not_use_matmul
self.modules_dtype_dict = modules_dtype_dict
@@ -1001,7 +1012,7 @@ class SDNQConfig(QuantizationConfigMixin):
value = list(value)
self.modules_dtype_dict[key] = value
if not isinstance(key, str) or not isinstance(value, list):
- raise ValueError(f"modules_dtype_dict must be a dictionary of strings and lists but got {type(key)} and {type(value)}")
+ raise TypeError(f"modules_dtype_dict must be a dictionary of strings and lists but got {type(key)} and {type(value)}")
if self.modules_quant_config is None:
self.modules_quant_config = {}
@@ -1027,18 +1038,23 @@ class SDNQConfig(QuantizationConfigMixin):
return f"SDNQConfig(weights_dtype={self.weights_dtype} quantization_device={self.quantization_device} return_device={self.return_device} group_size={self.group_size} use_quantized_matmul={self.use_quantized_matmul} quantized_matmul_dtype={self.quantized_matmul_dtype} quant_conv={self.quant_conv} quant_embedding={self.quant_embedding} use_quantized_matmul_conv={self.use_quantized_matmul_conv} use_static_quantization={self.use_static_quantization} use_dynamic_quantization={self.use_dynamic_quantization} dynamic_loss_threshold={self.dynamic_loss_threshold} use_stochastic_rounding={self.use_stochastic_rounding} use_hadamard={self.use_hadamard} hadamard_group_size={self.hadamard_group_size} use_svd={self.use_svd} svd_rank={self.svd_rank} svd_steps={self.svd_steps} dequantize_fp32={self.dequantize_fp32} non_blocking={self.non_blocking} add_skip_keys={self.add_skip_keys} modules_to_not_convert={self.modules_to_not_convert} modules_to_not_use_matmul={self.modules_to_not_use_matmul} modules_dtype_dict={self.modules_dtype_dict} modules_quant_config={self.modules_quant_config} )"
-import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position
-diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
-diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+if (
+ os.environ.get("SDNQ_REGISTER_DIFFUSERS", "0").lower() not in {"0", "false", "no"}
+ or (diffusers_version[0] == 0 and diffusers_version[1] < 40)
+):
+ import diffusers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position
+ diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+ diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+ diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
+ diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
-diffusers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
-diffusers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
-import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position
-transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
-transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+if os.environ.get("SDNQ_REGISTER_TRANSFORMERS", "1").lower() not in {"0", "false", "no"}:
+ import transformers.quantizers.auto # noqa: E402,RUF100 # pylint: disable=wrong-import-order,wrong-import-position
+ transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq"] = SDNQQuantizer
+ transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq"] = SDNQConfig
+ transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
+ transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
-transformers.quantizers.auto.AUTO_QUANTIZER_MAPPING["sdnq_training"] = SDNQQuantizer
-transformers.quantizers.auto.AUTO_QUANTIZATION_CONFIG_MAPPING["sdnq_training"] = SDNQConfig
sdnq_quantize_layer_weight_compiled = compile_func(sdnq_quantize_layer_weight)
diff --git a/modules/sdnq/utils.py b/modules/sdnq/utils.py
index cc423432f..bdc865087 100644
--- a/modules/sdnq/utils.py
+++ b/modules/sdnq/utils.py
@@ -44,13 +44,23 @@ def check_param_name_in(param_name: str, param_list: list[str]) -> str:
def check_quant_is_allowed(layer_class_name: str, weight: torch.Tensor, quantization_config, pre_quantized: bool = False) -> bool:
- return bool(
+ if (
layer_class_name in allowed_types
and weight.dtype in {torch.float64, torch.float32, torch.float16, torch.bfloat16}
and not (layer_class_name in embedding_types and not quantization_config.quant_embedding)
and not ((layer_class_name in conv_types or layer_class_name in conv_transpose_types) and not quantization_config.quant_conv)
- and (pre_quantized or weight.numel() >= quantization_config.minimum_allowed_numel)
- )
+ ):
+ if pre_quantized:
+ return True
+ if layer_class_name in conv_types:
+ channel_size = weight.shape[1]
+ elif layer_class_name in conv_transpose_types:
+ channel_size = weight.shape[0]
+ else:
+ channel_size = weight.shape[-1]
+ if channel_size >= quantization_config.minimum_allowed_channel_size and weight.numel() >= quantization_config.minimum_allowed_numel:
+ return True
+ return False
def check_quantized_matmul_is_allowed(use_quantized_matmul: bool, output_channel_size: int, channel_size: int) -> bool:
@@ -80,7 +90,7 @@ def get_quant_args_from_config(quantization_config: dict) -> dict:
quantization_config_dict.pop("is_training", None)
quantization_config_dict.pop("sdnq_version", None)
if quantization_config_dict.get("modules_quant_config", None) is not None:
- for key in quantization_config_dict["modules_quant_config"].keys():
+ for key in quantization_config_dict["modules_quant_config"]:
quantization_config_dict["modules_quant_config"][key] = get_quant_args_from_config(quantization_config_dict["modules_quant_config"][key])
return quantization_config_dict
@@ -90,7 +100,7 @@ def get_minimum_dtype(weights_dtype: str, param_name: str, modules_dtype_dict: d
for key, value in modules_dtype_dict.items():
if check_param_name_in(param_name, value) is not None:
key = key.lower()
- if key.startswith("minimum") or key.endswith("bit") or key.endswith("bits"):
+ if key.startswith("minimum") or key.endswith(("bit", "bits")):
minimum_bits_str = key.removeprefix("minimum").removeprefix("-").removeprefix("_").removesuffix("bits").removesuffix("bit").removesuffix("-").removesuffix("_")
if minimum_bits_str.startswith("uint"):
is_unsigned = True
@@ -189,7 +199,7 @@ def add_module_skip_keys(model: torch.nn.Module, quantization_config):
if skip_key_list is not None:
quantization_config.modules_to_not_convert.extend(skip_key_list[0])
for key, value in skip_key_list[1].items():
- if key in quantization_config.modules_dtype_dict.keys():
+ if key in quantization_config.modules_dtype_dict:
quantization_config.modules_dtype_dict[key].extend(value)
else:
quantization_config.modules_dtype_dict[key] = value
diff --git a/modules/seedvr/src/common/decorators.py b/modules/seedvr/src/common/decorators.py
deleted file mode 100644
index 52ab8ac03..000000000
--- a/modules/seedvr/src/common/decorators.py
+++ /dev/null
@@ -1,126 +0,0 @@
-import functools
-import threading
-from typing import Callable
-import torch
-from .distributed import barrier_if_distributed, get_global_rank, get_local_rank
-from .logger import get_logger
-
-
-logger = get_logger(__name__)
-
-
-def log_on_entry(func: Callable) -> Callable:
- """
- Functions with this decorator will log the function name at entry.
- When using multiple decorators, this must be applied innermost to properly capture the name.
- """
-
- def log_on_entry_wrapper(*args, **kwargs):
- logger.info(f"Entering {func.__name__}")
- return func(*args, **kwargs)
-
- return log_on_entry_wrapper
-
-
-def barrier_on_entry(func: Callable) -> Callable:
- """
- Functions with this decorator will start executing when all ranks are ready to enter.
- """
-
- def barrier_on_entry_wrapper(*args, **kwargs):
- barrier_if_distributed()
- return func(*args, **kwargs)
-
- return barrier_on_entry_wrapper
-
-
-def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable:
- """
- Helper function for local_rank_zero_only and global_rank_zero_only.
- """
-
- def conditional_execute_wrapper(*args, **kwargs):
- # Only execute if needed.
- result = func(*args, **kwargs) if execute else None
- # All GPUs must wait.
- barrier_if_distributed()
- # Return results.
- return result
-
- return conditional_execute_wrapper
-
-
-def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable:
- """
- Helper function for some functions with special constraints,
- especially functions called by other global_rank_zero_only / local_rank_zero_only ones,
- in case they are wrongly invoked in other scenarios.
- """
-
- def asserted_execute_wrapper(*args, **kwargs):
- assert condition, err_msg
- result = func(*args, **kwargs)
- return result
-
- return asserted_execute_wrapper
-
-
-def local_rank_zero_only(func: Callable) -> Callable:
- """
- Functions with this decorator will only execute on local rank zero.
- """
- return _conditional_execute_wrapper_factory(get_local_rank() == 0, func)
-
-
-def global_rank_zero_only(func: Callable) -> Callable:
- """
- Functions with this decorator will only execute on global rank zero.
- """
- return _conditional_execute_wrapper_factory(get_global_rank() == 0, func)
-
-
-def assert_only_global_rank_zero(func: Callable) -> Callable:
- """
- Functions with this decorator are only accessible to processes with global rank zero.
- """
- return _asserted_wrapper_factory(
- get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0"
- )
-
-
-def assert_only_local_rank_zero(func: Callable) -> Callable:
- """
- Functions with this decorator are only accessible to processes with local rank zero.
- """
- return _asserted_wrapper_factory(
- get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0"
- )
-
-
-def new_thread(func: Callable) -> Callable:
- """
- Functions with this decorator will run in a new thread.
- The function will return the thread, which can be joined to wait for completion.
- """
-
- def new_thread_wrapper(*args, **kwargs):
- thread = threading.Thread(target=func, args=args, kwargs=kwargs)
- thread.start()
- return thread
-
- return new_thread_wrapper
-
-
-def log_runtime(func: Callable) -> Callable:
- """
- Functions with this decorator will logging the runtime.
- """
-
- @functools.wraps(func)
- def wrapped(*args, **kwargs):
- barrier_if_distributed()
- result = func(*args, **kwargs)
- barrier_if_distributed()
- return result
-
- return wrapped
diff --git a/modules/seedvr/src/common/distributed/ops.py b/modules/seedvr/src/common/distributed/ops.py
index 593a5d427..ac01bbdab 100644
--- a/modules/seedvr/src/common/distributed/ops.py
+++ b/modules/seedvr/src/common/distributed/ops.py
@@ -357,7 +357,7 @@ def _broadcast_data(data, shape, dtype, src, group, async_op):
return comms
-def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]:
+def _traverse(data: Any, op: Callable) -> Union[List, Dict, Any, None]:
if isinstance(data, (list, tuple)):
return [_traverse(sub_data, op) for sub_data in data]
elif isinstance(data, dict):
diff --git a/modules/seedvr/src/common/seed.py b/modules/seedvr/src/common/seed.py
index 2469ad944..1129ebd5f 100644
--- a/modules/seedvr/src/common/seed.py
+++ b/modules/seedvr/src/common/seed.py
@@ -16,13 +16,14 @@ import random
from typing import Optional
import numpy as np
import torch
-from .distributed import get_global_rank
-def set_seed(seed: Optional[int], same_across_ranks: bool = False):
+def set_seed(seed: Optional[int]):
"""Function that sets the seed for pseudo-random number generators."""
+ if (seed is None) or (seed == '') or (seed == -1):
+ random.seed()
+ seed = int(random.randrange(4294967294))
if seed is not None:
- seed += get_global_rank() if not same_across_ranks else 0
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
diff --git a/modules/seedvr/src/core/generation.py b/modules/seedvr/src/core/generation.py
index 0b1d42f0f..811bdac88 100644
--- a/modules/seedvr/src/core/generation.py
+++ b/modules/seedvr/src/core/generation.py
@@ -107,7 +107,19 @@ def cut_videos(videos):
return result
-def generation_loop(runner, images, cfg_scale=1.0, cfg_rescale=0.0, steps=1, seed=666, res_w=720, batch_size=90, temporal_overlap=0, progress_callback=None, device:str='cpu', color_reconstruct=True):
+def generation_loop(runner,
+ images,
+ cfg_scale=1.5,
+ cfg_rescale=0.0,
+ steps=1,
+ seed=-1,
+ res_w=720,
+ batch_size=1,
+ temporal_overlap=0,
+ progress_callback=None,
+ device:str='cpu',
+ color_reconstruct=True,
+ ):
"""
Main generation loop with context-aware temporal processing
diff --git a/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py b/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
index 78a2d294d..39cc26dc4 100644
--- a/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
+++ b/modules/seedvr/src/models/video_vae_v3/modules/attn_video_vae.py
@@ -114,26 +114,14 @@ class Upsample3D(Upsample2D):
hidden_states = [hidden_states]
# ADD BY NUMZ
for i in range(len(hidden_states)):
- if self.use_conv and hasattr(self, "upscale_conv") and self.upscale_conv.kernel_size == (1, 1, 1):
- hidden_states[i] = hidden_states[i].repeat_interleave(self.temporal_ratio, dim=2)
- if self.spatial_ratio != 1:
- hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=3)
- hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=4)
- elif self.use_conv:
- hidden_states[i] = self.upscale_conv(hidden_states[i])
- hidden_states[i] = rearrange(
- hidden_states[i],
- "b (x y z c) f h w -> b c (f z) (h x) (w y)",
- x=self.spatial_ratio,
- y=self.spatial_ratio,
- z=self.temporal_ratio,
- ).contiguous()
- else:
- if self.temporal_ratio != 1:
- hidden_states[i] = hidden_states[i].repeat_interleave(self.temporal_ratio, dim=2)
- if self.spatial_ratio != 1:
- hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=3)
- hidden_states[i] = hidden_states[i].repeat_interleave(self.spatial_ratio, dim=4)
+ hidden_states[i] = self.upscale_conv(hidden_states[i])
+ hidden_states[i] = rearrange(
+ hidden_states[i],
+ "b (x y z c) f h w -> b c (f z) (h x) (w y)",
+ x=self.spatial_ratio,
+ y=self.spatial_ratio,
+ z=self.temporal_ratio,
+ )
# [Overridden] For causal temporal conv
if self.temporal_up and memory_state != MemoryState.ACTIVE:
@@ -1165,10 +1153,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
else:
encoded = self._encode(x)
posterior = DiagonalGaussianDistribution(encoded)
-
if not return_dict:
return (posterior,)
-
return AutoencoderKLOutput(latent_dist=posterior)
@apply_forward_hook
@@ -1181,16 +1167,15 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
decoded = self.tiled_decode(z)
else:
decoded = self._decode(z)
-
if not return_dict:
return (decoded,)
-
return DecoderOutput(sample=decoded)
def _encode(
self, x: torch.Tensor, memory_state: MemoryState = MemoryState.DISABLED
) -> torch.Tensor:
- _x = causal_conv_slice_inputs(x.to(self.device), self.slicing_sample_min_size, memory_state=memory_state)
+ _x = x.to(self.device)
+ _x = causal_conv_slice_inputs(_x, self.slicing_sample_min_size, memory_state=memory_state)
h = self.encoder(_x, memory_state=memory_state)
if self.quant_conv is not None:
output = self.quant_conv(h, memory_state=memory_state)
@@ -1262,109 +1247,53 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
overlap_size = int(self.tile_sample_min_size * (1 - self.tile_overlap_factor))
blend_extent = int(self.tile_latent_min_size * self.tile_overlap_factor)
row_limit = self.tile_latent_min_size - blend_extent
- prev_row = None
+ rows = []
self.tiles = 0
-
- row_positions = list(range(0, x.shape[3], overlap_size))
- col_positions = list(range(0, x.shape[4], overlap_size))
- enc = None
- output_width = 0
- h_cursor = 0
-
- for _row_idx, i in enumerate(row_positions):
- row_tiles = []
- for tile_idx, j in enumerate(col_positions):
+ for i in range(0, x.shape[3], overlap_size):
+ row = []
+ for j in range(0, x.shape[4], overlap_size):
tile = x[:, :, :, i : i + self.tile_sample_min_size, j : j + self.tile_sample_min_size]
tile = self._encode(tile)
- if tile.ndim == 4:
- tile = tile.unsqueeze(0)
- if prev_row is not None:
- tile = self.blend_v(prev_row[tile_idx], tile, blend_extent)
- if tile_idx > 0:
- tile = self.blend_h(row_tiles[-1], tile, blend_extent)
- row_tiles.append(tile)
+ row.append(tile)
self.tiles += 1
-
- cropped_tiles = [tile[:, :, :, :row_limit, :row_limit] for tile in row_tiles]
- row_width = 0
- for cropped in cropped_tiles:
- row_width += cropped.shape[-1]
- if output_width < cropped.shape[-1]:
- output_width = cropped.shape[-1]
-
- if enc is None:
- enc = torch.empty(
- cropped_tiles[0].shape[0],
- cropped_tiles[0].shape[1],
- cropped_tiles[0].shape[2],
- len(row_positions) * row_limit,
- len(col_positions) * row_limit,
- dtype=cropped_tiles[0].dtype,
- device=cropped_tiles[0].device,
- )
-
- w_cursor = 0
- for cropped in cropped_tiles:
- enc[:, :, :, h_cursor : h_cursor + cropped.shape[-2], w_cursor : w_cursor + cropped.shape[-1]] = cropped
- w_cursor += cropped.shape[-1]
-
- h_cursor += cropped_tiles[0].shape[-2]
- prev_row = row_tiles
-
- return enc[:, :, :, :h_cursor, :w_cursor]
+ rows.append(row)
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_extent)
+ result_row.append(tile[:, :, :, :row_limit, :row_limit])
+ result_rows.append(torch.cat(result_row, dim=4))
+ enc = torch.cat(result_rows, dim=3)
+ return enc
def tiled_decode(self, z: torch.Tensor) -> torch.Tensor:
overlap_size = int(self.tile_latent_min_size * (1 - self.tile_overlap_factor))
blend_extent = int(self.tile_sample_min_size * self.tile_overlap_factor)
row_limit = self.tile_sample_min_size - blend_extent
- prev_row = None
-
- row_positions = list(range(0, z.shape[3], overlap_size))
- col_positions = list(range(0, z.shape[4], overlap_size))
- dec = None
- output_width = 0
- h_cursor = 0
-
- for _row_idx, i in enumerate(row_positions):
- row_tiles = []
- for tile_idx, j in enumerate(col_positions):
+ rows = []
+ for i in range(0, z.shape[3], overlap_size):
+ row = []
+ for j in range(0, z.shape[4], overlap_size):
tile = z[:, :, :, i : i + self.tile_latent_min_size, j : j + self.tile_latent_min_size]
decoded = self.decoder(tile)
- if decoded.ndim == 4:
- decoded = decoded.unsqueeze(0)
- if prev_row is not None:
- decoded = self.blend_v(prev_row[tile_idx], decoded, blend_extent)
- if tile_idx > 0:
- decoded = self.blend_h(row_tiles[-1], decoded, blend_extent)
- row_tiles.append(decoded)
-
- cropped_tiles = [tile[:, :, :, :row_limit, :row_limit] for tile in row_tiles]
- row_width = 0
- for cropped in cropped_tiles:
- row_width += cropped.shape[-1]
- if output_width < cropped.shape[-1]:
- output_width = cropped.shape[-1]
-
- if dec is None:
- dec = torch.empty(
- cropped_tiles[0].shape[0],
- cropped_tiles[0].shape[1],
- cropped_tiles[0].shape[2],
- len(row_positions) * row_limit,
- len(col_positions) * row_limit,
- dtype=cropped_tiles[0].dtype,
- device=cropped_tiles[0].device,
- )
-
- w_cursor = 0
- for cropped in cropped_tiles:
- dec[:, :, :, h_cursor : h_cursor + cropped.shape[-2], w_cursor : w_cursor + cropped.shape[-1]] = cropped
- w_cursor += cropped.shape[-1]
-
- h_cursor += cropped_tiles[0].shape[-2]
- prev_row = row_tiles
-
- return dec[:, :, :, :h_cursor, :w_cursor]
+ row.append(decoded)
+ rows.append(row)
+ result_rows = []
+ for i, row in enumerate(rows):
+ result_row = []
+ for j, tile in enumerate(row):
+ if i > 0:
+ tile = self.blend_v(rows[i - 1][j], tile, blend_extent)
+ if j > 0:
+ tile = self.blend_h(row[j - 1], tile, blend_extent)
+ result_row.append(tile[:, :, :, :row_limit, :row_limit])
+ result_rows.append(torch.cat(result_row, dim=4))
+ dec = torch.cat(result_rows, dim=3)
+ return dec
def forward(
self, x: torch.FloatTensor, mode: Literal["encode", "decode", "all"] = "all", **kwargs
@@ -1405,6 +1334,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
):
self.spatial_downsample_factor = spatial_downsample_factor
self.temporal_downsample_factor = temporal_downsample_factor
+ self.freeze_encoder = freeze_encoder
self.freeze_encoder = True
super().__init__(*args, **kwargs)
diff --git a/modules/server.py b/modules/server.py
index ae17d541c..e00f40e56 100644
--- a/modules/server.py
+++ b/modules/server.py
@@ -1,36 +1,59 @@
import threading
import logging
+import time
+import asyncio
import uvicorn
import fastapi
+from modules.logger import log
class UvicornServer(uvicorn.Server):
- def __init__(self, app: fastapi.FastAPI, listen = None, port = None, keyfile = None, certfile = None, loop = "auto", http = "auto"):
+ def __init__(self, app: fastapi.FastAPI, host = None, listen = None, port = None, keyfile = None, certfile = None, loop = "auto", http = "auto"):
self.app: fastapi.FastAPI = app
self.thread: threading.Thread = None
+ self.loop = None
self.wants_restart = False
self.should_exit = False
+ kwargs = {
+ 'loop': loop, # auto, asyncio, uvloop
+ 'http': http, # auto, h11, httptools
+ 'interface': "auto", # auto, asgi3, asgi2, wsgi
+ 'ws': "auto", # auto, websockets, wsproto, websockets-sansio
+ 'timeout_keep_alive': 60, # default=5
+ 'ws_max_size': 1024 * 1024 * 1024, # default 16MB
+ 'ws_max_queue': 64, # default=32
+ 'ws_ping_interval': 30, # default=20
+ 'ws_ping_timeout': 60, # default=20
+ 'timeout_graceful_shutdown': 5, # default=None
+ 'access_log': False, # default=True
+ 'server_header': False, # default=True
+ 'date_header': False, # default=True
+ 'backlog': 4096, # default=2048
+ 'reload': False, # default=False
+ }
self.config = uvicorn.Config(
app=self.app,
- host = "0.0.0.0" if listen else "127.0.0.1",
- port = port or 7861,
- loop = loop, # auto, asyncio, uvloop
- http = http, # auto, h11, httptools
- interface = "auto", # auto, asgi3, asgi2, wsgi
- ws = "auto", # auto, websockets, wsproto
+ host = host or ("0.0.0.0" if listen else "127.0.0.1"),
+ port = port or 7860,
log_level = logging.WARNING,
- backlog = 4096, # default=2048
- timeout_keep_alive = 60, # default=5
ssl_keyfile = keyfile,
ssl_certfile = certfile,
- ws_max_size = 1024 * 1024 * 1024, # default 16MB
+ **kwargs
)
super().__init__(config=self.config)
+ log.info(f'Server: uvicorn={kwargs}')
def start(self):
self.thread = threading.Thread(target=self.run, daemon=True)
self.wants_restart = False
self.thread.start()
+ start = time.time()
+ while not self.started:
+ time.sleep(1e-3)
+ if time.time() - start > 5:
+ raise RuntimeError("Server failed to start. Please check that the port is available.")
+ policy = asyncio.get_event_loop_policy()
+ self.loop = f"{type(policy).__module__}.{type(policy).__name__}"
def stop(self):
self.should_exit = True
@@ -71,7 +94,6 @@ class HypercornServer:
self.thread = threading.Thread(target=self.run, daemon=True)
self.thread.start()
elif self.loop == 'asyncio': # does not run in thread
- import asyncio
from hypercorn.asyncio import serve
self.server = serve(self.app, self.config)
asyncio.run(self.server)
diff --git a/modules/shared_defaults.py b/modules/shared_defaults.py
index b680424cf..a20216823 100644
--- a/modules/shared_defaults.py
+++ b/modules/shared_defaults.py
@@ -1,3 +1,4 @@
+import os
from modules.logger import log
from modules import devices
@@ -58,6 +59,9 @@ def get_default_modes(cmd_opts, mem_stat):
elif devices.backend in {"directml", "cpu", "mps"}:
default_sdp_override_options = ['Dynamic attention']
+ if devices.get_optimal_device_name() != "cpu":
+ os.environ.setdefault('SDNQ_USE_OPENVINO_MM', '0')
+
return (
default_offload_mode,
default_diffusers_offload_min_gpu_memory,
diff --git a/modules/shared_items.py b/modules/shared_items.py
index b2dbb9deb..a69edfeba 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -84,6 +84,8 @@ pipelines = {
'XOmni': None,
'ZetaChroma': None,
'Boogu': None,
+ 'SeFi': None,
+ 'MageFlow': None,
}
diff --git a/modules/shared_state.py b/modules/shared_state.py
index 8acc8a86b..34fd0172f 100644
--- a/modules/shared_state.py
+++ b/modules/shared_state.py
@@ -293,8 +293,9 @@ class State:
elif self.prediction_type == "v_prediction":
sample = self.current_noise_pred * (-self.current_sigma / (self.current_sigma**2 + 1) ** 0.5) + (original_sample / (self.current_sigma**2 + 1)) # pylint: disable=invalid-unary-operand-type
except Exception:
+ # log.error(f'State image sigma: last={self.id_live_preview} step={self.sampling_step} {e}')
pass # ignore sigma errors
- image = sd_samplers_common.samples_to_image_grid(sample)
+ image = sd_samplers_common.samples_to_image_grid(sample, fast=self.sampling_step > 1)
self.assign_current_image(image)
self.preview_job = -1
return True
diff --git a/modules/sharpfin/sparse_backend.py b/modules/sharpfin/sparse_backend.py
index 87c673ef1..a20fcc1dd 100644
--- a/modules/sharpfin/sparse_backend.py
+++ b/modules/sharpfin/sparse_backend.py
@@ -462,7 +462,7 @@ def triton_dds(
fuse_srgb: bool = False,
clamp_output: bool = False,
output_mt: bool = False,
- output_slice: None | Tuple[int,int] = None
+ output_slice: Tuple[int, int] | None = None
):
assert isinstance(lhs, torch.Tensor)
assert isinstance(rhs, Matrix)
@@ -616,7 +616,7 @@ def triton_dds_sbsc(
fuse_srgb: bool = False,
clamp_output: bool = False,
output_mt: bool = False,
- output_slice: None | Tuple[int,int] = None
+ output_slice: Tuple[int, int] | None = None
):
assert isinstance(lhs, torch.Tensor)
assert isinstance(rhs, SBSCMatrix)
@@ -776,7 +776,7 @@ def triton_dds_zerorhs_sbsc(
gamma_correction: str = 'fast',
clamp_output: bool = False,
output_mt: bool = False,
- output_slice: None | Tuple[int,int] = None
+ output_slice: Tuple[int, int] | None = None
):
assert isinstance(lhs, torch.Tensor)
diff --git a/modules/storage.py b/modules/storage.py
new file mode 100644
index 000000000..b41a7e7c7
--- /dev/null
+++ b/modules/storage.py
@@ -0,0 +1,176 @@
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+from modules import paths
+from modules.logger import log
+from modules.shared import opts, max_workers
+from modules.modelstats import stat
+
+
+class Location:
+ name: str
+ folders: list[str]
+ paths: list[Path]
+ nfiles: int = 0
+ nfolders: int = 0
+ nsymlinks: int = 0
+ nerrors: int = 0
+ type: str = ''
+ size: int = 0
+ time: float = 0.0
+ mtime: datetime = datetime.fromtimestamp(0)
+
+ def __init__(self, name: str | None, folders: str | list[str], what: str = ''):
+ self.type = what
+ if isinstance(folders, str):
+ folders = [folders]
+ self.name = name if name is not None else ', '.join(folders)
+ self.folders = folders
+ self.paths = [Path(f).resolve(strict=False) for f in self.folders if f is not None and f != '']
+
+ def __repr__(self):
+ return f'Location(type={self.type} name="{self.name}" folders={self.folders} size={self.size/1024/1024:.3f} mtime="{self.mtime}" files={self.nfiles} folders={self.nfolders} symlinks={self.nsymlinks} errors={self.nerrors} time={self.time:.3f})'
+
+ def dict(self):
+ return {
+ 'name': self.name,
+ 'type': self.type,
+ 'folders': self.folders,
+ 'paths': [str(p) for p in self.paths],
+ 'size': self.size,
+ 'mtime': self.mtime.timestamp(),
+ 'nfiles': self.nfiles,
+ 'nfolders': self.nfolders,
+ 'nsymlinks': self.nsymlinks,
+ 'nerrors': self.nerrors,
+ 'time': self.time,
+ }
+
+def get_all_locations(types: list[str] | None = []) -> list[Location]:
+ locations = []
+ if types is None or 'All' in types or 'Models' in types:
+ locations.append(Location('SD Models', opts.ckpt_dir, 'Models'))
+ locations.append(Location('Diffusers Models', opts.diffusers_dir, 'Models'))
+ locations.append(Location('Huggingface Modules', opts.hfcache_dir, 'Models'))
+ locations.append(Location('VAE', [opts.vae_dir, os.path.join(paths.models_path, "TAESD")], 'Models'))
+ locations.append(Location('UNet', opts.unet_dir, 'Models'))
+ locations.append(Location('TextEncoder', opts.te_dir, 'Models'))
+ locations.append(Location('LoRA', opts.lora_dir, 'Models'))
+ locations.append(Location('ControlNets', opts.control_dir, 'Models'))
+ locations.append(Location('Embeddings', opts.embeddings_dir, 'Models'))
+ locations.append(Location('Detailers', [opts.yolo_dir, os.path.join(paths.models_path, 'Ultralytics')], 'Models'))
+ locations.append(Location('Upscalers', [opts.esrgan_models_path, opts.bsrgan_models_path, opts.realesrgan_models_path, opts.scunet_models_path, opts.swinir_models_path, os.path.join(paths.models_path, 'chaiNNer'), os.path.join(paths.models_path, 'GFPGAN'), os.path.join(paths.models_path, 'Spandrel'), os.path.join(paths.models_path, 'SeedVR2')], 'Models')) # chainners extension has late opts init
+ locations.append(Location('CLiP', opts.clip_models_path, 'Models'))
+ locations.append(Location('Rembg', os.path.join(paths.models_path, 'Rembg'), 'Models'))
+ locations.append(Location('RIFE', os.path.join(paths.models_path, 'RIFE'), 'Models'))
+ if types is None or 'All' in types or 'Data' in types:
+ locations.append(Location('Configs', ['data', paths.sd_configs_path], 'Data'))
+ locations.append(Location('AutoComplete', opts.autocomplete_dir, 'Data'))
+ locations.append(Location('Styles', opts.styles_dir, 'Data'))
+ locations.append(Location('Wildcards', opts.wildcards_dir, 'Data'))
+ locations.append(Location('Reference', paths.reference_path, 'Data'))
+ locations.append(Location('LUTs', os.path.join(paths.models_path, 'LUTs'), 'Data'))
+ locations.append(Location('Wiki', 'wiki', 'Data'))
+ if types is None or 'All' in types or 'Cache' in types:
+ locations.append(Location('Temp', opts.temp_dir, 'Cache'))
+ locations.append(Location('XET', opts.xetcache_dir, 'Cache'))
+ locations.append(Location('OpenVINO', opts.openvino_cache_path, 'Cache'))
+ locations.append(Location('ONNX', opts.onnx_cached_models_path, 'Cache'))
+ locations.append(Location('VENV', 'venv', 'Cache'))
+ locations.append(Location('Torch', [opts.tunable_dir, os.getenv("TORCHINDUCTOR_CACHE_DIR", None), os.getenv("TRITON_CACHE_DIR", None)], 'Cache'))
+ if types is None or 'All' in types or 'Code' in types:
+ locations.append(Location('Modules', 'modules', 'Code'))
+ locations.append(Location('Pipelines', 'pipelines', 'Code'))
+ locations.append(Location('Scripts', 'scripts', 'Code'))
+ locations.append(Location('UI', 'ui', 'Code'))
+ locations.append(Location('Builtin', paths.extensions_builtin_dir, 'Code'))
+ locations.append(Location('Extensions', paths.extensions_dir, 'Code'))
+ if types is None or 'All' in types or 'Images' in types:
+ locations.append(Location('Text', [opts.outdir_txt2img_samples], 'Images'))
+ locations.append(Location('Image', [opts.outdir_img2img_samples], 'Images'))
+ locations.append(Location('Control', [opts.outdir_control_samples], 'Images'))
+ locations.append(Location('Extras', [opts.outdir_extras_samples], 'Images'))
+ locations.append(Location('Save', [opts.outdir_save], 'Images'))
+ locations.append(Location('Grids', [opts.outdir_txt2img_grids, opts.outdir_img2img_grids, opts.outdir_control_grids], 'Images'))
+ if types is None or 'All' in types or 'Videos' in types:
+ locations.append(Location('Video', [opts.outdir_video], 'Videos'))
+ return locations
+
+
+def get_other_locations(locations: list[Location], name: str, folder: str, what: str = 'Other') -> list[Location]:
+ # get list of first level subfolders in `folder` check each if its already in `locations` by comparing resolved paths if not, add each to the list as a new Location with type `what`
+ existing_paths = set()
+ for location in locations:
+ for path in location.paths:
+ existing_paths.add(path.resolve(strict=False))
+ try:
+ with os.scandir(folder) as entries:
+ for entry in entries:
+ if entry.is_dir(follow_symlinks=False):
+ path = Path(entry.path).resolve(strict=False)
+ if path not in existing_paths:
+ locations.append(Location(name, entry.path, what))
+ except (FileNotFoundError, PermissionError):
+ pass
+ return locations
+
+
+def print_summary(locations: list[Location]):
+ summary = {}
+ for location in locations:
+ if location.type not in summary:
+ summary[location.type] = {
+ 'size': 0,
+ 'mtime': datetime.fromtimestamp(0),
+ 'nfiles': 0,
+ 'nfolders': 0,
+ 'nerrors': 0,
+ }
+ summary[location.type]['size'] += location.size
+ if location.mtime > summary[location.type]['mtime']:
+ summary[location.type]['mtime'] = location.mtime
+ summary[location.type]['nfiles'] += location.nfiles
+ summary[location.type]['nfolders'] += location.nfolders
+ summary[location.type]['nerrors'] += location.nerrors
+ for k, v in summary.items():
+ log.debug(f'Storage: type={k} size={v["size"]/1024/1024:.3f} files={v["nfiles"]} folders={v["nfolders"]}')
+
+
+def check_storage(folders: str | list[str] | None = None, types: list[str] | None = None, silent: bool = False) -> list[Location]:
+ if isinstance(folders, str):
+ folders = [folders]
+ if folders is not None and len(folders) > 0:
+ locations = [Location(None, folder) for _i, folder in enumerate(folders)]
+ else:
+ locations = get_all_locations(types)
+ if types is None or 'Other' in types or 'All' in types:
+ locations = get_other_locations(locations, 'Other', paths.models_path)
+ log.debug(f'Storage: locations={len(locations)} workers={max_workers} types={types} folders={folders} start')
+
+ def update_stats(location: Location) -> Location:
+ t0 = time.time()
+ for f in location.paths:
+ size, mtime, files, folders, symlinks, errors = stat(f, extended=True, exclude=['__pycache__', '.'])
+ location.size += size
+ if mtime > location.mtime:
+ location.mtime = mtime
+ location.nfiles += files
+ location.nfolders += folders
+ location.nsymlinks += symlinks
+ location.nerrors += errors
+ location.time = time.time() - t0
+ return location
+
+ t0 = time.time()
+ from concurrent.futures import ThreadPoolExecutor, as_completed
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
+ future_items = {executor.submit(update_stats, location): location for location in locations}
+ for future in as_completed(future_items):
+ location = future.result()
+ if location.size > 0 and not silent:
+ log.debug(location)
+ t1 = time.time()
+ print_summary(locations)
+ log.debug(f'Storage: time={t1-t0:.3f} end')
+ return locations
diff --git a/modules/ui_common.py b/modules/ui_common.py
index 78583ea9f..ad2612868 100644
--- a/modules/ui_common.py
+++ b/modules/ui_common.py
@@ -296,7 +296,7 @@ def open_folder(result_gallery, gallery_index = 0):
subprocess.Popen([opener, path]) # pylint: disable=consider-using-with
-def create_output_panel(tabname, preview=True, prompt=None, height=None, transfer=True, scale=1, result_info=None):
+def create_output_panel(tabname, preview=True, prompt=None, height=None, transfer=True, scale=1, result_info=None, html_log_val=''):
with gr.Column(variant='panel', elem_id=f"{tabname}_results", scale=scale):
with gr.Group(elem_id=f"{tabname}_gallery_container"):
if tabname == "txt2img":
@@ -348,7 +348,7 @@ def create_output_panel(tabname, preview=True, prompt=None, height=None, transfe
html_info = gr.HTML(elem_id=f'html_info_{tabname}', elem_classes="infotext", visible=False) # contains raw infotext as returned by wrapped call
html_info_formatted = gr.HTML(elem_id=f'html_info_formatted_{tabname}', elem_classes="infotext", visible=True) # contains html formatted infotext
html_info.change(fn=infotext_to_html, inputs=[html_info], outputs=[html_info_formatted], show_progress='hidden')
- html_log = gr.HTML(elem_id=f'html_log_{tabname}', elem_classes=["hint"])
+ html_log = gr.HTML(elem_id=f'html_log_{tabname}', elem_classes=["hint"], value=html_log_val)
generation_info = gr.Textbox(visible=False, elem_id=f'generation_info_{tabname}')
generation_info_button = gr.Button(visible=False, elem_id=f"{tabname}_generation_info_button")
diff --git a/modules/ui_control.py b/modules/ui_control.py
index 37414d9df..951e0ba4b 100644
--- a/modules/ui_control.py
+++ b/modules/ui_control.py
@@ -21,7 +21,7 @@ use_generator = os.environ.get('SD_USE_GENERATOR', None) is not None
def return_stats(t: float | None = None):
if t is None:
- elapsed_text = ''
+ elapsed_text = '⏱ Idle'
else:
elapsed = time.perf_counter() - t
elapsed_m = int(elapsed // 60)
@@ -35,15 +35,16 @@ def return_stats(t: float | None = None):
ooms = mem_mon_read.pop("oom")
retries = mem_mon_read.pop("retries")
vram = {k: v // 1048576 for k, v in mem_mon_read.items()}
- peak = max(vram['active_peak'], vram['reserved_peak'], vram['used'])
- used = round(100.0 * peak / vram['total']) if vram['total'] > 0 else 0
+ peak = max(vram.get('active_peak', 0), vram.get('reserved_peak', 0), vram.get('used', 0))
+ used = round(100.0 * peak / vram.get('total', 0)) if vram.get('total', 0) > 0 else 0
if peak > 0:
gpu += f"| 🕮 GPU {peak} MB"
gpu += f" {used}%" if used > 0 else ''
- gpu += f" | retries {retries} oom {ooms}" if retries > 0 or ooms > 0 else ''
+ gpu += f" | Retries {retries} OOM {ooms}" if retries > 0 or ooms > 0 else ''
ram = ram_stats()
if ram['used'] > 0:
- cpu += f" RAM {ram['used']} GB"
+ # change emoji/symbol for ram to something better
+ cpu += f"| 🗒 RAM {ram['used']} GB"
cpu += f" {round(100.0 * ram['used'] / ram['total'])}%" if ram['total'] > 0 else ''
return f""
@@ -247,7 +248,7 @@ def create_ui(_blocks: gr.Blocks=None):
gr.HTML('Output
')
with gr.Tabs(elem_classes=['control-tabs'], elem_id='control-tab-output') as output_tabs:
with gr.Tab('Gallery', id='out-gallery'):
- output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=False, prompt=prompt, height=gr_height, result_info=result_txt)
+ output_gallery, _output_gen_info, _output_html_info, _output_html_info_formatted, output_html_log = ui_common.create_output_panel("control", preview=False, prompt=prompt, height=gr_height, result_info=result_txt, html_log_val=return_stats())
with gr.Tab('Image', id='out-image'):
output_image = gr.Image(label="Output", show_label=False, type="pil", interactive=False, tool="editor", height=gr_height, elem_id='control_output_image', elem_classes=['control-image'])
with gr.Tab('Video', id='out-video'):
diff --git a/modules/ui_definitions.py b/modules/ui_definitions.py
index 2a38277dd..3e5d11572 100644
--- a/modules/ui_definitions.py
+++ b/modules/ui_definitions.py
@@ -150,7 +150,7 @@ def create_settings(cmd_opts):
"caption_offload": OptionInfo(True, "Offload caption models"),
"caption_to_gpu": OptionInfo(True, "Load caption models direct to GPU"),
"offload_balanced_sep": OptionInfo("Balanced Offload
", "", gr.HTML),
- "diffusers_offload_pre": OptionInfo(True, "Offload during pre-forward"),
+ "diffusers_offload_pre": OptionInfo(True, "Offload during pre-forward", gr.Checkbox, {"visible": False}),
"diffusers_offload_streams": OptionInfo(False, "Offload using streams"),
"diffusers_offload_min_gpu_memory": OptionInfo(startup_offload_min_gpu, "Offload low watermark", gr.Slider, {"minimum": 0, "maximum": 1, "step": 0.01 }),
"diffusers_offload_max_gpu_memory": OptionInfo(startup_offload_max_gpu, "Offload GPU high watermark", gr.Slider, {"minimum": 0.1, "maximum": 1, "step": 0.01 }),
diff --git a/modules/ui_extra_networks.py b/modules/ui_extra_networks.py
index 301edaa94..1315e275c 100644
--- a/modules/ui_extra_networks.py
+++ b/modules/ui_extra_networks.py
@@ -1,5 +1,6 @@
import os
import io
+from functools import lru_cache
import random
import re
import time
@@ -14,6 +15,7 @@ from html.parser import HTMLParser
from collections import OrderedDict
import gradio as gr
from PIL import Image
+from fastapi.exceptions import HTTPException
from starlette.responses import FileResponse, JSONResponse
from modules import paths, shared, devices, files_cache, errors, infotext, ui_symbols, ui_components, modelstats
from modules.logger import log
@@ -55,11 +57,12 @@ preview_map = None
def init_api():
def get_thumb(filename: str = ""):
- global allowed_dirs # pylint: disable=global-statement
- if len(allowed_dirs) == 0:
- allowed_dirs = shared.demo.allowed_paths
+ if os.path.join('ui', 'assets') not in allowed_dirs:
+ allowed_dirs.append(os.path.join('ui', 'assets'))
if filename is None or len(filename) == 0:
return JSONResponse({ "error": "no filename" }, status_code=400)
+ if not any(Path(folder).absolute() in Path(filename).absolute().parents for folder in allowed_dirs):
+ raise HTTPException(status_code=403, detail=f"file {filename}: must be in one of allowed directories")
if not os.path.exists(filename) or not os.path.isfile(filename) or os.path.getsize(filename) == 0:
return FileResponse('ui/assets/missing.png', headers={"Accept-Ranges": "bytes"})
if filename.startswith('html/') or filename.startswith('models/') or filename.startswith('data/') or filename.startswith('ui/'):
@@ -196,8 +199,12 @@ class ExtraNetworksPage:
errors.display(e, 'Network version')
return all_versions[0]
+ @lru_cache(maxsize=2048, typed=False)
def link_preview(self, filename: str):
- if not os.path.exists(filename):
+ if filename == 'ui/assets/missing.png':
+ return f"{shared.opts.subpath}/sdapi/v1/network/thumb?filename={filename}"
+ just_file = not bool(os.path.dirname(filename))
+ if just_file or not os.path.exists(filename):
ref = os.path.join(paths.reference_path, filename)
if os.path.exists(ref):
filename = ref
@@ -456,10 +463,13 @@ class ExtraNetworksPage:
errors.display(e, 'Networks')
return ""
+ @lru_cache(maxsize=2048, typed=False)
def find_preview_file(self, path: str | None):
if path is None:
return 'ui/assets/missing.png'
if os.path.join('models', 'Reference') in path:
+ if shared.cmd_opts.test and not os.path.exists(path):
+ log.warning(f'Networks: missing-preview type="{self.name}" fn="{path}"')
return path
exts = ["jpg", "jpeg", "png", "webp", "tiff", "jp2", "jxl"]
reference_path = os.path.abspath(os.path.join('models', 'Reference'))
@@ -477,6 +487,7 @@ class ExtraNetworksPage:
return file
return 'ui/assets/missing.png'
+ @lru_cache(maxsize=2048, typed=False)
def find_preview(self, filename: str):
t0 = time.time()
preview_file = self.find_preview_file(filename)
diff --git a/modules/ui_extra_networks_checkpoints.py b/modules/ui_extra_networks_checkpoints.py
index 93a6e2f09..a68fb0daa 100644
--- a/modules/ui_extra_networks_checkpoints.py
+++ b/modules/ui_extra_networks_checkpoints.py
@@ -63,6 +63,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
v['tags'].append(f'Size: {size} GB')
shared.reference_models[k] = v
+ models = []
for k, v in shared.reference_models.items():
count['total'] += 1
url = v['path']
@@ -126,14 +127,14 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
if ready:
count['ready'] += 1
- yield {
+ model = {
"type": 'Model',
"name": name,
"title": name,
"filename": url,
"preview": self.find_preview(os.path.join(paths.reference_path, preview)),
"local_preview": preview_file,
- "onclick": '"' + html.escape(f"selectReference({json.dumps(path)})") + '"',
+ "onclick": '"' + html.escape(f"""return selectReference({json.dumps(path)})""") + '"',
"hash": None,
"mtime": mtime,
"size": size,
@@ -143,7 +144,9 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
"version": version,
"tags": v.get('tags', []),
}
+ models.append(model)
log.debug(f'Networks: type="reference" {count}')
+ return models
def create_item(self, name):
record = None
@@ -157,7 +160,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
"filename": checkpoint.filename,
"hash": checkpoint.shorthash,
"metadata": checkpoint.metadata,
- "onclick": '"' + html.escape(f"selectCheckpoint({json.dumps(name)})") + '"',
+ "onclick": '"' + html.escape(f"""return selectCheckpoint({json.dumps(name)})""") + '"',
"mtime": mtime,
"size": size,
}
@@ -190,8 +193,7 @@ class ExtraNetworksPageCheckpoints(ui_extra_networks.ExtraNetworksPage):
item = future.result()
if item is not None:
items.append(item)
- for record in self.list_reference():
- items.append(record)
+ items += self.list_reference()
self.update_all_previews(items)
return items
diff --git a/modules/ui_extra_networks_history.py b/modules/ui_extra_networks_history.py
index 3d2d05059..3dbd61829 100644
--- a/modules/ui_extra_networks_history.py
+++ b/modules/ui_extra_networks_history.py
@@ -19,9 +19,10 @@ class ExtraNetworksPageHistory(ui_extra_networks.ExtraNetworksPage):
def list_items(self):
# log.trace('History list')
+ records = []
for item in shared.history.latents:
title = ', '.join(list(set(item.ops))) + '
' + item.name
- yield {
+ record = {
"type": 'History',
"name": title,
"preview": item.preview,
@@ -31,6 +32,8 @@ class ExtraNetworksPageHistory(ui_extra_networks.ExtraNetworksPage):
# "description": item.info,
"onclick": '"' + html.escape(f"""return selectHistory({json.dumps(item.name)})""") + '"',
}
+ records.append(record)
+ return records
def find_description(self, path, info=None):
name = path.split('
')[-1]
diff --git a/modules/ui_extra_networks_styles.py b/modules/ui_extra_networks_styles.py
index 793e6d836..5eca82c0a 100644
--- a/modules/ui_extra_networks_styles.py
+++ b/modules/ui_extra_networks_styles.py
@@ -137,6 +137,5 @@ class ExtraNetworkStyles(extra_networks.ExtraNetwork):
p.negative_prompts = [styles.merge_prompts(style.negative_prompt, prompt) for prompt in p.negative_prompts]
styles.apply_styles_to_extra(p, style)
-
def deactivate(self, p, force=False):
pass
diff --git a/modules/ui_extra_networks_unet.py b/modules/ui_extra_networks_unet.py
index 6b03b7be0..7badcdfc0 100644
--- a/modules/ui_extra_networks_unet.py
+++ b/modules/ui_extra_networks_unet.py
@@ -13,6 +13,7 @@ class ExtraNetworksPageUNets(ui_extra_networks.ExtraNetworksPage):
return sd_unet.refresh_unet_list()
def list_items(self):
+ results = []
for name, filename in sd_unet.unet_dict.items():
try:
size, mtime = modelstats.stat(filename)
@@ -35,9 +36,10 @@ class ExtraNetworksPageUNets(ui_extra_networks.ExtraNetworksPage):
"description": self.find_description(filename, info),
"version": version.get("baseModel", "N/A") if info else "N/A",
}
- yield record
+ results.append(record)
except Exception as e:
log.debug(f'Networks error: type=unet file="{filename}" {e}')
+ return results
def allowed_directories_for_previews(self):
return [v for v in [shared.opts.unet_dir] if v is not None]
diff --git a/modules/ui_extra_networks_vae.py b/modules/ui_extra_networks_vae.py
index 20a890544..b3341c621 100644
--- a/modules/ui_extra_networks_vae.py
+++ b/modules/ui_extra_networks_vae.py
@@ -10,9 +10,10 @@ class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage):
super().__init__('VAE')
def refresh(self):
- shared.refresh_vaes()
+ return shared.refresh_vaes()
def list_items(self):
+ records = []
for name, filename in sd_vae.vae_dict.items():
try:
size, mtime = modelstats.stat(filename)
@@ -35,9 +36,10 @@ class ExtraNetworksPageVAEs(ui_extra_networks.ExtraNetworksPage):
"description": self.find_description(filename, info),
"version": version.get("baseModel", "N/A") if info else "N/A",
}
- yield record
+ records.append(record)
except Exception as e:
log.debug(f'Networks error: type=vae file="{filename}" {e}')
+ return records
def allowed_directories_for_previews(self):
return [v for v in [shared.opts.vae_dir] if v is not None]
diff --git a/modules/ui_extra_networks_wildcards.py b/modules/ui_extra_networks_wildcards.py
index 8538eaeb4..489053124 100644
--- a/modules/ui_extra_networks_wildcards.py
+++ b/modules/ui_extra_networks_wildcards.py
@@ -40,6 +40,7 @@ class ExtraNetworksPageWildcards(ui_extra_networks.ExtraNetworksPage):
def list_items(self):
self.refresh()
+ records = []
for filename in wildcards_list:
relname = os.path.relpath(filename, shared.opts.wildcards_dir)
name = os.path.splitext(relname)[0]
@@ -49,17 +50,20 @@ class ExtraNetworksPageWildcards(ui_extra_networks.ExtraNetworksPage):
"type": 'Wildcard',
"name": name,
"filename": filename,
- "preview": self.find_preview(filename),
+ # "preview": self.find_preview(filename),
+ "preview": None,
"local_preview": f"{os.path.splitext(filename)[0]}.{shared.opts.samples_format}",
"prompt": json.dumps(f" __{name}__"),
+ # "prompt": f" __{name}__",
"mtime": mtime,
"size": size,
"description": '',
"info": {},
}
- yield record
+ records.append(record)
except Exception as e:
log.debug(f'Networks error: type=wildcard file="{filename}" {e}')
+ return records
def allowed_directories_for_previews(self):
return [v for v in [shared.opts.wildcards_dir] if v is not None]
diff --git a/modules/ui_settings.py b/modules/ui_settings.py
index 547273f2e..74787c431 100644
--- a/modules/ui_settings.py
+++ b/modules/ui_settings.py
@@ -1,7 +1,7 @@
import os
import gradio as gr
from modules import timer, shared, paths, theme, sd_models, modelloader, generation_parameters_copypaste, call_queue, script_callbacks
-from modules import ui_common, ui_loadsave, ui_history, ui_components, ui_symbols
+from modules import ui_common, ui_loadsave, ui_history, ui_storage, ui_components, ui_symbols
from modules.logger import log
@@ -313,6 +313,10 @@ def create_ui(disabled_tabs=None):
with gr.TabItem("History", id="system_history", elem_id="tab_history"):
ui_history.create_ui()
+ if 'storage' not in disabled_tabs:
+ with gr.TabItem("Storage", id="system_storage", elem_id="tab_storage"):
+ ui_storage.create_ui()
+
if 'monitor' not in disabled_tabs:
with gr.TabItem("GPU Monitor", id="system_gpu", elem_id="tab_gpu"):
with gr.Row(elem_id='gpu-controls'):
diff --git a/modules/ui_storage.py b/modules/ui_storage.py
new file mode 100644
index 000000000..dd4952549
--- /dev/null
+++ b/modules/ui_storage.py
@@ -0,0 +1,13 @@
+import gradio as gr
+
+
+def create_ui():
+ types = ['All', 'Images', 'Videos', 'Models', 'Data', 'Cache', 'Code', 'Other']
+ with gr.Row():
+ btn_refresh = gr.Button("Calculate", elem_id='btn_storage_refresh')
+ storage_type = gr.Dropdown(label="Storage type", elem_id='storage_type', choices=types, value=[types[0]], multiselect=True)
+ with gr.Row():
+ _storage_table = gr.HTML('', elem_id='storage_table')
+ with gr.Row():
+ _storage_timeline = gr.HTML('', elem_id='storage_timeline')
+ btn_refresh.click(_js='refreshStorage', fn=None, inputs=[storage_type], outputs=[], show_progress='full')
diff --git a/modules/ui_video.py b/modules/ui_video.py
index 910befcc8..c87c516a8 100644
--- a/modules/ui_video.py
+++ b/modules/ui_video.py
@@ -1,6 +1,6 @@
import os
import gradio as gr
-from modules import shared, timer, images, ui_common, ui_sections, generation_parameters_copypaste
+from modules import shared, timer, images, ui_common, ui_sections, generation_parameters_copypaste, scripts_manager
from modules.logger import log
@@ -9,6 +9,10 @@ debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lam
def create_ui():
log.debug('UI initialize: tab=video')
+
+ scripts_manager.scripts_current = scripts_manager.scripts_video
+ scripts_manager.scripts_video.initialize_scripts(is_img2img=False, is_control=False, is_video=True)
+
with gr.Blocks(analytics_enabled=False) as _video_interface:
prompt, styles, negative, generate_btn, _reprocess, paste, networks_button, _token_counter, _token_button, _token_counter_negative, _token_button_negative = ui_sections.create_toprow(
is_img2img=False,
@@ -31,25 +35,31 @@ def create_ui():
with gr.Tab('Output', id='video-outputs-tab') as _video_outputs_tab:
from modules.video_models import video_ui
mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb = video_ui.create_ui_outputs()
+ with gr.Tab('Extras', id='video-extras-tab', elem_id='video_extras') as _video_extras_tab:
+ video_script_inputs = scripts_manager.scripts_video.setup_ui(parent='video', accordion=True)
with gr.Tab('Generic', id='video-core-tab') as video_core_tab:
from modules.video_models import video_ui
engine, model, steps, sampler_index, width, height, frames, seed = video_ui.create_ui(
- prompt, negative, styles, overrides,
+ prompt, negative, styles,
+ overrides, video_script_inputs,
mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb,
)
with gr.Tab('FramePack', id='framepack-tab') as framepack_tab:
from modules.framepack import framepack_ui
framepack_ui.create_ui(
- prompt, negative, styles, overrides,
+ prompt, negative, styles,
+ overrides, video_script_inputs,
mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb,
)
with gr.Tab('LTX', id='ltx-tab') as ltx_tab:
from modules.ltx import ltx_ui
ltx_ui.create_ui(
- prompt, negative, styles, overrides,
+ prompt, negative, styles,
+ overrides, video_script_inputs,
mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb,
)
+
paste_fields = [
(prompt, "Prompt"), # cannot add more fields as they are not defined yet
(negative, "Negative prompt"),
@@ -75,5 +85,3 @@ def create_ui():
ltx_tab.select(fn=lambda: 'ltx', inputs=[], outputs=[current_tab])
generate_btn.click(fn=None, _js='submit_video_wrapper', inputs=[current_tab], outputs=[])
-
- # from framepack_api import create_api # pylint: disable=wrong-import-order
diff --git a/modules/ui_video_vlm.py b/modules/ui_video_vlm.py
index da3ed1ec7..9646d8375 100644
--- a/modules/ui_video_vlm.py
+++ b/modules/ui_video_vlm.py
@@ -1,3 +1,6 @@
+# legacy module as video now uses main prompt enhancer
+# except for framepack
+
import gradio as gr
from PIL import Image
from modules.logger import log
diff --git a/modules/upscaler.py b/modules/upscaler.py
index 6bafa3db0..b02633f58 100644
--- a/modules/upscaler.py
+++ b/modules/upscaler.py
@@ -118,7 +118,7 @@ class Upscaler:
img = self.do_upscale(img, selected_model)
if shape == (img.width, img.height):
break
- if img.width >= dest_w and img.height >= dest_h:
+ if img.width >= (dest_w - 8) and img.height >= (dest_h - 8):
break
if img.width != dest_w or img.height != dest_h:
from modules.image import sharpfin
diff --git a/modules/vae/sd_vae_taesd.py b/modules/vae/sd_vae_taesd.py
index 00739f13f..5ea9eac9c 100644
--- a/modules/vae/sd_vae_taesd.py
+++ b/modules/vae/sd_vae_taesd.py
@@ -43,7 +43,8 @@ prev_warnings = False
first_run = True
prev_cls = ''
prev_type = ''
-prev_model = ''
+prev_variant = ''
+prev_model = None
lock = threading.Lock()
@@ -82,12 +83,12 @@ def get_model(model_cls, variant=None):
warn_once(f'cls={shared.sd_model.__class__.__name__} type={shared.sd_model_type} unsuppported', variant=variant)
return model_cls, None
if debug:
- log.debug(f'TAESD detect: cls={model_cls} variant={variant}')
+ log.debug(f'TAESD detect: cls={model_cls} variant="{variant}"')
return model_cls, variant
def load_model(model_type = 'decoder', variant = None, vae_file: str | None = None):
- global prev_cls, prev_type, prev_model, prev_warnings # pylint: disable=global-statement
+ global prev_cls, prev_type, prev_variant, prev_warnings # pylint: disable=global-statement
model_cls = shared.sd_model_type if shared.sd_loaded else None
if vae_file is not None and os.path.exists(vae_file):
model_cls = 'sdxl'
@@ -101,7 +102,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
os.makedirs(folder, exist_ok=True)
if variant.startswith('TAE'):
cfg = TAESD_MODELS[variant]
- if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+ if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_variant) and (cfg['model'] is not None):
return cfg['model'], variant
fn = os.path.join(folder, cfg['fn'] + model_type + '_' + model_cls + '.pth')
if not os.path.exists(fn):
@@ -117,7 +118,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
if os.path.exists(fn):
prev_cls = model_cls
prev_type = model_type
- prev_model = variant
+ prev_variant = variant
log.print() # new line
log.debug(f'Decode: type="taesd" variant="{variant}" fn="{fn}" layers={shared.opts.taesd_layers} load')
vae = None
@@ -133,12 +134,6 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
else:
from modules.taesd.taesd import TAESD
vae = TAESD(decoder_path=fn if model_type=='decoder' else None, encoder_path=fn if model_type=='encoder' else None)
- """
- _vae = diffusers.AutoencoderKL()
- from installer import Dot
- _config = diffusers.AutoencoderKL().config.copy()
- vae.config = Dot(_config) # set config for compatibility with standard vae
- """
if vae is not None:
prev_warnings = False # reset warnings for new model
vae = vae.to(devices.device, dtype=dtype)
@@ -147,7 +142,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
return vae, variant
elif variant.startswith('Hybrid'):
cfg = CQYAN_MODELS[variant].get(model_cls, None)
- if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_model) and (cfg['model'] is not None):
+ if (model_cls == prev_cls) and (model_type == prev_type) and (variant == prev_variant) and (cfg['model'] is not None):
return cfg['model'], variant
if cfg is None:
warn_once(f'cls={shared.sd_model.__class__.__name__} type={model_cls} unsuppported', variant=variant)
@@ -155,7 +150,7 @@ def load_model(model_type = 'decoder', variant = None, vae_file: str | None = No
repo = cfg['repo']
prev_cls = model_cls
prev_type = model_type
- prev_model = variant
+ prev_variant = variant
log.debug(f'Decode: type="taesd" variant="{variant}" id="{repo}" load')
if 'tiny' in repo:
from diffusers.models import AutoencoderTiny
@@ -190,13 +185,20 @@ def restore_preview_size(image, vae):
return image
-def decode(latents):
- global first_run # pylint: disable=global-statement
+def decode(latents, fast=False):
+ global first_run, prev_model, prev_variant # pylint: disable=global-statement
with lock:
try:
- vae, variant = load_model(model_type='decoder')
- if vae is None or max(latents.shape) > 256: # safetey check of large tensors
- return latents
+ if fast and prev_model is not None:
+ vae = prev_model
+ variant = prev_variant
+ else:
+ vae, variant = load_model(model_type='decoder')
+ if vae is None or max(latents.shape) > 256: # safety check of large tensors
+ return latents
+ prev_model = vae
+ prev_variant = variant
+ fast = False
except Exception as e:
# from modules import errors
# errors.display(e, 'taesd"')
@@ -208,7 +210,7 @@ def decode(latents):
tensor = latents.unsqueeze(0) if len(latents.shape) == 3 else latents
tensor = tensor.detach().clone().to(devices.device, dtype=dtype)
if debug:
- log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} tensor={tensor.shape}')
+ log.debug(f'Decode: type="taesd" variant="{variant}" input={latents.shape} fast={fast} tensor={tensor.shape}')
# Fallback: reshape packed 128-channel latents to 32 channels if not already unpacked
if (variant == 'TAE FLUX.2') and (len(tensor.shape) == 4) and (tensor.shape[1] == 128):
b, _c, h, w = tensor.shape
@@ -221,7 +223,7 @@ def decode(latents):
image = (image / 2.0 + 0.5).clamp(0, 1).detach()
image = restore_preview_size(image, vae)
t1 = time.time()
- if (t1 - t0) > 3.0 and not first_run:
+ if (t1 - t0) > 5.0 and not first_run:
log.warning(f'Decode: type="taesd" variant="{variant}" long decode time={t1 - t0:.2f}')
first_run = False
return image
diff --git a/modules/video_models/video_prompt.py b/modules/video_models/video_prompt.py
deleted file mode 100644
index 963813e9d..000000000
--- a/modules/video_models/video_prompt.py
+++ /dev/null
@@ -1,24 +0,0 @@
-from modules import shared, extra_networks, ui_video_vlm
-
-
-def prepare_prompts(p, init_image, prompt:str, vlm_enhance:bool, vlm_model:str, vlm_system_prompt:str):
- p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
- p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
- shared.prompt_styles.apply_styles_to_extra(p)
- p.prompts, p.network_data = extra_networks.parse_prompts([p.prompt])
- extra_networks.activate(p)
- prompt = p.prompts[0]
-
- new_prompt = ui_video_vlm.enhance_prompt(
- enable=vlm_enhance,
- model=vlm_model,
- image=init_image,
- prompt=prompt,
- system_prompt=vlm_system_prompt,
- )
- if new_prompt is not None and len(new_prompt) > 0:
- prompt = new_prompt
-
- p.styles = []
- p.task_args['prompt'] = p.prompt
- p.task_args['negative_prompt'] = p.negative_prompt
diff --git a/modules/video_models/video_run.py b/modules/video_models/video_run.py
index 93a666480..195543c85 100644
--- a/modules/video_models/video_run.py
+++ b/modules/video_models/video_run.py
@@ -1,17 +1,27 @@
import os
import copy
import time
-from modules import shared, errors, sd_models, processing, devices, images, ui_common
+from modules import shared, errors, sd_models, processing, devices, images, ui_common, scripts_manager
from modules.logger import log
-from modules.video_models import models_def, video_utils, video_load, video_vae, video_overrides, video_save, video_prompt
+from modules.video_models import models_def, video_utils, video_load, video_vae, video_overrides, video_save
from modules.paths import resolve_output_path
debug = log.trace if os.environ.get('SD_VIDEO_DEBUG', None) is not None else lambda *args, **kwargs: None
-def generate(*args, **kwargs):
- task_id, ui_state, engine, model, prompt, negative, styles, width, height, frames, steps, sampler_index, sampler_shift, dynamic_shift, seed, guidance_scale, guidance_true, init_image, init_strength, last_image, vae_type, vae_tile_frames, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb, vlm_enhance, vlm_model, vlm_system_prompt, override_settings = args
+def generate(task_id, ui_state,
+ engine, model,
+ prompt, negative, styles,
+ width, height, frames, steps,
+ sampler_index, sampler_shift, dynamic_shift,
+ seed, guidance_scale, guidance_true,
+ init_image, init_strength, last_image,
+ vae_type, vae_tile_frames,
+ mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb,
+ override_settings,
+ *args, **kwargs
+ ):
if engine is None or model is None or engine == 'None' or model == 'None':
return video_utils.queue_err('model not selected')
@@ -54,9 +64,12 @@ def generate(*args, **kwargs):
if p.vae_type == 'Remote' and not selected.vae_remote:
log.warning(f'Video: model={selected.name} remote vae not supported')
p.vae_type = 'Default'
- p.scripts = None
- p.script_args = None
+
p.state = ui_state
+ p.scripts = scripts_manager.scripts_video
+ p.script_args = args
+ processed: processing.Processed = scripts_manager.scripts_video.run(p, *args)
+
p.do_not_save_grid = True
p.do_not_save_samples = not mp4_frames
p.outpath_samples = resolve_output_path(shared.opts.outdir_samples, shared.opts.outdir_video)
@@ -101,9 +114,7 @@ def generate(*args, **kwargs):
shared.sd_model = sd_models.apply_balanced_offload(shared.sd_model)
devices.torch_gc(force=True, reason='video')
-
# set args
- video_prompt.prepare_prompts(p, init_image, prompt, vlm_enhance, vlm_model, vlm_system_prompt)
processing.fix_seed(p)
video_vae.set_vae_params(p)
p.task_args['num_inference_steps'] = p.steps
diff --git a/modules/video_models/video_save.py b/modules/video_models/video_save.py
index 49166e237..b14274ee6 100644
--- a/modules/video_models/video_save.py
+++ b/modules/video_models/video_save.py
@@ -103,55 +103,62 @@ def numpy_to_tensor(images):
return tensor
-def add_audio_stream(container, audio_sample_rate: int):
- # Must be registered before the first container.mux(); avformat_write_header runs there
- # and freezes the stream set, after which new streams have time_base=0/0.
- audio_stream = container.add_stream("aac", rate=audio_sample_rate)
- audio_stream.codec_context.sample_rate = audio_sample_rate
- audio_stream.codec_context.layout = "stereo"
- audio_stream.codec_context.time_base = Fraction(1, audio_sample_rate)
- log.debug(f'Audio: codec={audio_stream.codec_context.name} rate={audio_stream.codec_context.sample_rate} layout={audio_stream.codec_context.layout} base={audio_stream.codec_context.time_base}')
- return audio_stream
+def add_audio_packets(container, audio_stream, audio: dict):
+ if not audio or "frames" not in audio:
+ return
+ try:
+ av = check_av()
+ sr = audio.get("sr", 44100)
+ layout = audio.get("layout", "stereo")
+ resampler = av.AudioResampler(format="fltp", layout=layout, rate=sr)
+ fifo = av.AudioFifo()
+ for raw_frame in audio.get("frames", []):
+ for resampled in resampler.resample(raw_frame):
+ fifo.write(resampled)
+ for resampled in resampler.resample(None):
+ fifo.write(resampled)
+ pts_counter = 0
+ frame_size = audio_stream.codec_context.frame_size or 1024
+ while fifo.samples >= frame_size:
+ frame = fifo.read(frame_size)
+ frame.pts = pts_counter
+ pts_counter += frame.samples
+ for packet in audio_stream.encode(frame):
+ packet.stream = audio_stream
+ container.mux_one(packet)
+ if fifo.samples > 0:
+ frame = fifo.read(fifo.samples)
+ frame.pts = pts_counter
+ pts_counter += frame.samples
+ for packet in audio_stream.encode(frame):
+ packet.stream = audio_stream
+ container.mux_one(packet)
+ for packet in audio_stream.encode():
+ packet.stream = audio_stream
+ container.mux_one(packet)
+ except Exception as e:
+ log.error(f"Video audio encoding: type=packets {e}")
+ errors.display(e, "Audio")
-def write_audio(
- container,
- audio_stream,
- samples: torch.Tensor,
- audio_sample_rate: int,
-) -> None:
+def add_audio_tensor(container, audio_stream, audio: torch.Tensor, sample_rate: int):
av = check_av()
- audio_stream.codec_context.format = "fltp"
- if samples.ndim == 1:
- samples = samples[:, None]
- if samples.shape[1] != 2 and samples.shape[0] == 2:
- samples = samples.T
- if samples.shape[1] != 2:
- raise ValueError(f"Expected samples with 2 channels; got shape {samples.shape}.")
- if samples.dtype != torch.int16:
- samples = torch.clip(samples, -1.0, 1.0)
- samples = (samples * 32767.0).to(torch.int16)
- audio_frames = av.AudioFrame.from_ndarray(
- samples.contiguous().reshape(1, -1).cpu().numpy(),
- format="s16",
- layout="stereo",
- )
- audio_frames.sample_rate = audio_sample_rate
- audio_resampler = av.audio.resampler.AudioResampler(
- format=audio_stream.codec_context.format,
- layout=audio_stream.codec_context.layout,
- rate=audio_stream.codec_context.sample_rate,
- )
- pts = 0
- for resampled in audio_resampler.resample(audio_frames):
- resampled.pts = resampled.pts or 0
- resampled.sample_rate = audio_frames.sample_rate
- packets = audio_stream.encode(resampled)
- for packet in packets:
- container.mux(packet)
- pts += resampled.samples
- for packet in audio_stream.encode():
- container.mux(packet)
+ if torch.is_tensor(audio):
+ audio = audio.detach().float().cpu().numpy()
+ if audio.ndim > 2:
+ audio = np.squeeze(audio)
+ if audio.ndim == 1:
+ audio = audio[None, :]
+ elif audio.ndim == 2 and audio.shape[0] > audio.shape[1] and audio.shape[1] in (1, 2):
+ audio = audio.T
+ channels = audio.shape[0] if audio.shape[0] in (1, 2) else 1
+ layout = "stereo" if channels == 2 else "mono"
+ if audio.dtype != np.int16:
+ audio = np.clip(audio, -1.0, 1.0)
+ audio = (audio * 32767.0).astype(np.int16)
+ audio_frame = av.AudioFrame.from_ndarray(audio, format="s16p", layout=layout)
+ audio_frame.sample_rate = sample_rate
+ add_audio_packets(container, audio_stream, {"sr": sample_rate, "layout": layout, "frames": [audio_frame]})
def atomic_save_video(
@@ -162,7 +169,7 @@ def atomic_save_video(
codec: str = "libx264",
pix_fmt: str = "yuv420p",
options: str = "",
- aac: int = 24000,
+ sample_rate: int = 24000,
metadata: dict | None = None,
pbar=None,
):
@@ -172,23 +179,23 @@ def atomic_save_video(
if av is None or av is False:
log.error('Video: ffmpeg/av not available')
return
-
savejob = shared.state.begin('Save video')
frames, height, width, _channels = tensor.shape
rate = round(fps)
- options_str = options
- options = {}
- for option in [option.strip() for option in options_str.split(',')]:
- if '=' in option:
- key, value = option.split('=', 1)
- elif ':' in option:
- key, value = option.split(':', 1)
- else:
- continue
- options[key.strip()] = value.strip()
- log.info(f'Video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} audio={audio is not None} aac={aac} options={options}')
+ parsed_options = {}
+ if isinstance(options, str):
+ for option in [opt.strip() for opt in options.split(',')]:
+ if '=' in option:
+ key, value = option.split('=', 1)
+ elif ':' in option:
+ key, value = option.split(':', 1)
+ else:
+ continue
+ parsed_options[key.strip()] = value.strip()
+ elif isinstance(options, dict):
+ parsed_options = options
+ log.info(f'Video: file="{filename}" codec={codec} frames={frames} width={width} height={height} fps={rate} audio={audio is not None} sample_rate={sample_rate} options={parsed_options}')
video_array = torch.as_tensor(tensor, dtype=torch.uint8).numpy(force=True)
-
task = pbar.add_task('encoding', total=frames) if pbar is not None else None
if task is not None:
pbar.update(task, description='video encoding')
@@ -196,25 +203,32 @@ def atomic_save_video(
with av.open(filename, mode="w") as container:
for k, v in metadata.items():
container.metadata[k] = v
- stream: av.VideoStream = container.add_stream(codec, rate=rate, options=options)
+ stream: av.VideoStream = container.add_stream(codec, rate=rate, options=parsed_options)
stream.width = video_array.shape[2]
stream.height = video_array.shape[1]
stream.pix_fmt = pix_fmt
- audio_stream = add_audio_stream(container, aac) if audio is not None else None
- for img in video_array:
+ stream.time_base = Fraction(1, rate)
+ audio_stream = None
+ has_audio = (audio is not None) and ((torch.is_tensor(audio) or isinstance(audio, np.ndarray)) or (isinstance(audio, dict) and len(audio.get('frames', [])) > 0))
+ if has_audio:
+ sr = sample_rate if not isinstance(audio, dict) else audio.get("sr", sample_rate)
+ layout = "stereo" if not isinstance(audio, dict) else audio.get("layout", "stereo")
+ audio_stream = container.add_stream("aac", rate=sr)
+ audio_stream.layout = layout
+ audio_stream.time_base = Fraction(1, sr)
+ for i, img in enumerate(video_array):
frame = av.VideoFrame.from_ndarray(img, format="rgb24")
+ frame.pts = i
for packet in stream.encode_lazy(frame):
- container.mux(packet)
+ container.mux_one(packet)
if task is not None:
pbar.update(task, advance=1)
- for packet in stream.encode(): # flush
- container.mux(packet)
- if audio_stream is not None:
- try:
- write_audio(container, audio_stream, audio, aac)
- except Exception as e:
- log.error(f'Video audio encoding: {e}')
- errors.display(e, 'Audio')
+ if (audio is not None) and (torch.is_tensor(audio) or isinstance(audio, np.ndarray)):
+ add_audio_tensor(container, audio_stream, audio, sample_rate)
+ elif (audio is not None) and isinstance(audio, dict) and len(audio.get('frames', [])) > 0:
+ add_audio_packets(container, audio_stream, audio)
+ for packet in stream.encode():
+ container.mux_one(packet)
shared.state.outputs(filename)
shared.state.end(savejob)
@@ -294,7 +308,13 @@ def save_video(
n, _c, t, h, w = pixels.shape
size = pixels.element_size() * pixels.numel()
log.debug(f'Video: video={mp4_video} export={mp4_frames} safetensors={mp4_sf} interpolate={mp4_interpolate}')
- log.debug(f'Video: encode={t} raw={size} latent={pixels.shape} audio={audio.shape if audio is not None else None} fps={mp4_fps} codec={mp4_codec} ext={mp4_ext} options="{mp4_opt}"')
+ if hasattr(audio, 'shape'):
+ audio_txt = f'audio={audio.shape} aac={aac_sample_rate}' if audio is not None else 'no audio'
+ elif isinstance(audio, dict):
+ audio_txt = f'audio={audio.get("format", None)} packets={len(audio.get("frames", []))} '
+ else:
+ audio_txt = None
+ log.debug(f'Video: encode={t} raw={size} latent={pixels.shape} {audio_txt} fps={mp4_fps} codec={mp4_codec} ext={mp4_ext} options="{mp4_opt}"')
try:
preparejob = shared.state.begin('Prepare video')
if stream is not None:
@@ -340,7 +360,7 @@ def save_video(
if mp4_video and (mp4_codec != 'none'):
output_video = f'{output_filename}.{mp4_ext}'
metadata = create_video_metadata(p, metadata, output_filename)
- atomic_save_video(output_video, tensor=x, audio=audio, fps=mp4_fps, codec=mp4_codec, options=mp4_opt, aac=aac_sample_rate, metadata=metadata, pbar=pbar)
+ atomic_save_video(output_video, tensor=x, audio=audio, fps=mp4_fps, codec=mp4_codec, options=mp4_opt, sample_rate=aac_sample_rate, metadata=metadata, pbar=pbar)
if stream is not None:
stream.output_queue.push(('progress', (None, f'Video {os.path.basename(output_video)} | Codec {mp4_codec} | Size {w}x{h}x{t} | FPS {mp4_fps}')))
stream.output_queue.push(('file', output_video))
diff --git a/modules/video_models/video_ui.py b/modules/video_models/video_ui.py
index 7b65111be..4bcc5b618 100644
--- a/modules/video_models/video_ui.py
+++ b/modules/video_models/video_ui.py
@@ -1,6 +1,6 @@
import os
import gradio as gr
-from modules import sd_models, ui_common, ui_sections, ui_symbols, ui_video_vlm, call_queue
+from modules import sd_models, ui_common, ui_sections, ui_symbols, call_queue
from modules.logger import log
from modules.ui_components import ToolButton
from modules.video_models import models_def, video_utils
@@ -105,7 +105,7 @@ def create_ui_outputs():
return mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb
-def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
+def create_ui(prompt, negative, styles, overrides, script_inputs, mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb):
with gr.Row():
with gr.Column(variant='compact', elem_id="video_settings", elem_classes=['settings-column'], scale=1):
with gr.Row():
@@ -143,8 +143,6 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
vae_type = gr.Dropdown(label='VAE decode', choices=['Default', 'Tiny', 'Remote', 'Upscale'], value='Default', elem_id="video_vae_type")
vae_tile_frames = gr.Slider(label='Tile frames', minimum=1, maximum=64, step=1, value=16, elem_id="video_vae_tile_frames")
- vlm_enhance, vlm_model, vlm_system_prompt = ui_video_vlm.create_ui(prompt_element=prompt, image_element=init_image)
-
# output panel with gallery and video tabs
with gr.Column(elem_id='video-output-column', scale=2) as _column_output:
with gr.Tabs(elem_classes=['video-output-tabs'], elem_id='video-output-tabs'):
@@ -174,7 +172,6 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
init_image, init_strength, last_image,
vae_type, vae_tile_frames,
mp4_fps, mp4_interpolate, mp4_codec, mp4_ext, mp4_opt, mp4_video, mp4_frames, mp4_sf, mp4_thumb,
- vlm_enhance, vlm_model, vlm_system_prompt,
overrides,
]
video_outputs = [
@@ -188,7 +185,7 @@ def create_ui(prompt, negative, styles, overrides, mp4_fps, mp4_interpolate, mp4
video_dict = dict(
fn=call_queue.wrap_gradio_gpu_call(video_run.generate, extra_outputs=[gr.update(), gr.update(), gr.update(), gr.update()], name='Video'),
_js="submit_video",
- inputs=state_inputs + video_inputs,
+ inputs=state_inputs + video_inputs + script_inputs,
outputs=video_outputs,
show_progress='hidden',
)
diff --git a/pipelines/anima/anima_image.py b/pipelines/anima/anima_image.py
index 5d0334181..c24fe8706 100644
--- a/pipelines/anima/anima_image.py
+++ b/pipelines/anima/anima_image.py
@@ -1,14 +1,12 @@
"""Anima img2img and inpainting pipelines (built dynamically from the runtime-imported base class)."""
from typing import Callable, Dict, List, Optional, Union
-
import torch
import torch.nn.functional as F
from PIL import Image
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import PipelineImageInput
from diffusers.utils.torch_utils import randn_tensor
-
from modules import devices
@@ -131,12 +129,20 @@ def build_anima_pipeline_classes(base_cls):
try:
return base_cls.__call__(
self,
- prompt=prompt, negative_prompt=negative_prompt,
- height=height, width=width, num_inference_steps=num_inference_steps,
- guidance_scale=guidance_scale, num_images_per_prompt=num_images_per_prompt,
- generator=generator, latents=noised, prompt_embeds=prompt_embeds,
- negative_prompt_embeds=negative_prompt_embeds, output_type=output_type,
- return_dict=return_dict, callback_on_step_end=blend_callback,
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ num_images_per_prompt=num_images_per_prompt,
+ generator=generator,
+ latents=noised,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ output_type=output_type,
+ return_dict=return_dict,
+ callback_on_step_end=blend_callback,
callback_on_step_end_tensor_inputs=["latents"],
max_sequence_length=max_sequence_length,
)
diff --git a/pipelines/bria/bria_pipeline.py b/pipelines/bria/bria_pipeline.py
index 9d5cd324d..a602dd05a 100644
--- a/pipelines/bria/bria_pipeline.py
+++ b/pipelines/bria/bria_pipeline.py
@@ -255,7 +255,7 @@ class BriaPipeline(FluxPipeline):
callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 128,
- clip_value:Union[None,float] = None,
+ clip_value:Union[float, None] = None,
normalize:bool = False
):
r"""
diff --git a/pipelines/generic_shared.py b/pipelines/generic_shared.py
index 1ad9ac3fc..554e0fa37 100644
--- a/pipelines/generic_shared.py
+++ b/pipelines/generic_shared.py
@@ -107,7 +107,7 @@ shared_te_map = {
'Qwen3-VL 4B Conditional': {
'cls': transformers.Qwen3VLForConditionalGeneration,
'target_repo': 'SeFi-Image/SeFi-Image-5B-Base',
- 'identifier': ['5b'],
+ 'identifier': ['4b','5b'],
'target_subfolder': 'Qwen3-VL-4B-Instruct',
},
'Qwen3-VL 8B Conditional': {
diff --git a/pipelines/ideogram/ideogram4.py b/pipelines/ideogram/ideogram4.py
index ed7b52533..d6ec40816 100644
--- a/pipelines/ideogram/ideogram4.py
+++ b/pipelines/ideogram/ideogram4.py
@@ -653,7 +653,7 @@ class Ideogram4Pipeline(DiffusionPipeline):
# 4. Set up the resolution-aware logit-normal schedule on the scheduler.
schedule_mu = _resolution_aware_mu(height=height, width=width, base_mu=mu)
sigmas = _logit_normal_sigmas(num_inference_steps, schedule_mu, std=std, device=device)
- self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device)
+ self.scheduler.set_timesteps(sigmas=sigmas.tolist(), device=device) # pylint: disable=unexpected-keyword-arg
timesteps = self.scheduler.timesteps
self._num_timesteps = len(timesteps) # pylint: disable=attribute-defined-outside-init
diff --git a/pipelines/krea2/pipeline_krea2_inpaint.py b/pipelines/krea2/pipeline_krea2_inpaint.py
new file mode 100644
index 000000000..30533a285
--- /dev/null
+++ b/pipelines/krea2/pipeline_krea2_inpaint.py
@@ -0,0 +1,132 @@
+"""Krea 2 inpainting pipeline.
+
+This module adds inpainting support on top of the existing Krea2 image-to-image
+variant without modifying the upstream Krea2 denoising path.
+"""
+
+from typing import Callable, Dict, List, Optional, Union
+
+import torch
+import torch.nn.functional as F
+from PIL import Image
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput
+from diffusers.utils.torch_utils import randn_tensor
+
+from modules import devices
+from .pipeline_krea2 import Krea2Pipeline, Krea2Img2ImgPipeline
+
+
+def _setup_img2img_schedule(scheduler, strength, num_inference_steps, device, mu=None):
+ """Set custom sigma schedule, return first sigma after scheduler shift."""
+ min_sigma = 1e-8
+ custom_sigmas = torch.linspace(max(strength, 0.01), min_sigma, num_inference_steps).tolist()
+ scheduler.set_timesteps(sigmas=custom_sigmas, device=device, mu=mu)
+ return scheduler.sigmas[0].item()
+
+
+def _prepare_mask(pipe, mask_image, height, width, device):
+ if isinstance(mask_image, list):
+ mask_image = mask_image[0]
+ if isinstance(mask_image, Image.Image):
+ mask_image = mask_image.convert("L")
+ if isinstance(mask_image, Image.Image):
+ import torchvision.transforms.functional as TF
+
+ mask_tensor = TF.to_tensor(mask_image).unsqueeze(0).to(device=device, dtype=torch.float32)
+ elif isinstance(mask_image, torch.Tensor):
+ mask_tensor = mask_image.to(device=device, dtype=torch.float32)
+ if mask_tensor.ndim == 2:
+ mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
+ elif mask_tensor.ndim == 3:
+ mask_tensor = mask_tensor.unsqueeze(0)
+ else:
+ mask_tensor = torch.ones(1, 1, height, width, device=device, dtype=torch.float32)
+
+ latent_h = height // pipe.vae_compression
+ latent_w = width // pipe.vae_compression
+ mask_latent = F.interpolate(mask_tensor, size=(latent_h, latent_w), mode="nearest")
+ return mask_latent[:, :1, :, :]
+
+
+class Krea2InpaintPipeline(Krea2Img2ImgPipeline):
+ """Krea 2 inpainting pipeline."""
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ image: Optional[PipelineImageInput] = None,
+ mask_image: Optional[PipelineImageInput] = None,
+ strength: float = 0.8,
+ height: int = 1024,
+ width: int = 1024,
+ num_inference_steps: int = 28,
+ guidance_scale: float | None = None,
+ num_images_per_prompt: int = 1,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.Tensor] = None,
+ output_type: str = "pil",
+ return_dict: bool = True,
+ callback_on_step_end: Optional[Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ ):
+ align = self.vae_compression * self.patch
+ height = (height // align) * align
+ width = (width // align) * align
+
+ device = devices.device
+ dtype = self.transformer.dtype
+
+ cfg = self.scheduler.config
+ grid_h = height // (self.vae_compression * self.patch)
+ grid_w = width // (self.vae_compression * self.patch)
+ mu = self.calculate_shift(
+ grid_h * grid_w,
+ cfg.get("base_image_seq_len", 256),
+ cfg.get("max_image_seq_len", 6400),
+ cfg.get("base_shift", 0.5),
+ cfg.get("max_shift", 1.15),
+ )
+ actual_sigma = _setup_img2img_schedule(self.scheduler, strength, num_inference_steps, device, mu=mu)
+ init_latents = self.encode_image(image, height, width, dtype, device)
+ noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=devices.dtype)
+ noised = (actual_sigma * noise + (1.0 - actual_sigma) * init_latents).to(torch.float32)
+ mask_latent = _prepare_mask(self, mask_image, height, width, device)
+
+ orig_set_timesteps = self.scheduler.set_timesteps
+ self.scheduler.set_timesteps = lambda *args, **kwargs: None
+
+ user_callback = callback_on_step_end
+
+ def blend_callback(pipe, i, t, callback_kwargs):
+ cur_latents = callback_kwargs.get("latents")
+ if cur_latents is not None:
+ sigma_next = pipe.scheduler.sigmas[i + 1].item() if i + 1 < len(pipe.scheduler.sigmas) else 0.0
+ init_at_t = sigma_next * noise + (1.0 - sigma_next) * init_latents
+ blended = mask_latent * cur_latents + (1.0 - mask_latent) * init_at_t.to(cur_latents.dtype)
+ callback_kwargs["latents"] = blended
+ if user_callback is not None:
+ callback_kwargs = user_callback(pipe, i, t, callback_kwargs)
+ return callback_kwargs
+
+ try:
+ return Krea2Pipeline.__call__(
+ self,
+ prompt=prompt,
+ negative_prompt=negative_prompt,
+ height=height,
+ width=width,
+ num_inference_steps=num_inference_steps,
+ guidance_scale=guidance_scale,
+ num_images_per_prompt=num_images_per_prompt,
+ generator=generator,
+ latents=noised,
+ output_type=output_type,
+ return_dict=return_dict,
+ callback_on_step_end=blend_callback,
+ callback_on_step_end_tensor_inputs=["latents"],
+ )
+ finally:
+ self.scheduler.set_timesteps = orig_set_timesteps
diff --git a/pipelines/mageflow/__init__.py b/pipelines/mageflow/__init__.py
new file mode 100644
index 000000000..137ebeeab
--- /dev/null
+++ b/pipelines/mageflow/__init__.py
@@ -0,0 +1,10 @@
+import diffusers
+from .pipeline_mage_flow import MageFlowPipeline
+from .pipeline_output import MageFlowPipelineOutput
+from .autoencoder_mage_vae import AutoencoderMageVAE
+from .transformer_mage_flow import MageFlowTransformer2DModel
+
+diffusers.MageFlowPipeline = MageFlowPipeline
+diffusers.MageFlowPipelineOutput = MageFlowPipelineOutput
+diffusers.AutoencoderMageVAE = AutoencoderMageVAE
+diffusers.MageFlowTransformer2DModel = MageFlowTransformer2DModel
diff --git a/pipelines/mageflow/autoencoder_mage_vae.py b/pipelines/mageflow/autoencoder_mage_vae.py
new file mode 100644
index 000000000..aa4f0ebc8
--- /dev/null
+++ b/pipelines/mageflow/autoencoder_mage_vae.py
@@ -0,0 +1,768 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+MageVAE: DConvEncoder + DConvDenoiser (with CoD Decoder) autoencoder.
+
+Encodes images to 128-channel latents at 16x spatial downsampling using a one-step
+diffusion encoder, and decodes latents back to images using a DConv denoiser conditioned
+on a CoD (Cascaded-of-Decoders) decoder.
+
+Latent shape: [B, 128, H/16, W/16].
+"""
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.autoencoders.vae import DecoderOutput
+
+
+logger = logging.get_logger(__name__)
+
+
+# ---------------------------------------------------------------------------
+# Helper
+# ---------------------------------------------------------------------------
+def _mage_vae_modulate(x, shift, scale):
+ if x.dim() == 4:
+ batch_size, channels = x.shape[:2]
+ return x * (1 + scale.view(batch_size, channels, 1, 1)) + shift.view(batch_size, channels, 1, 1)
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+# ---------------------------------------------------------------------------
+# Primitive layers
+# ---------------------------------------------------------------------------
+class MageVAELayerNorm2d(nn.LayerNorm):
+ """Channel-last LayerNorm for NCHW tensors."""
+
+ def __init__(self, num_channels, eps=1e-6, affine=True):
+ super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+ def forward(self, x):
+ x = x.permute(0, 2, 3, 1).contiguous()
+ x = F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+ return x.permute(0, 3, 1, 2).contiguous()
+
+
+class MageVAERMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, x):
+ input_dtype = x.dtype
+ x = x.to(torch.float32)
+ variance = x.pow(2).mean(-1, keepdim=True)
+ x = x * torch.rsqrt(variance + self.variance_epsilon)
+ return self.weight * x.to(input_dtype)
+
+
+class MageVAETimestepEmbedder(nn.Module):
+ """Timestep MLP (max_period=10000, freq_size=256)."""
+
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ half = dim // 2
+ freqs = torch.exp(-math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half).to(t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ embedding = self.timestep_embedding(t, self.frequency_embedding_size)
+ return self.mlp(embedding.to(self.mlp[0].weight.dtype))
+
+
+# ---------------------------------------------------------------------------
+# DConv blocks
+# ---------------------------------------------------------------------------
+class MageVAEDiCoBlock(nn.Module):
+ """DConv block with adaLN modulation, used in encoder and decoder."""
+
+ def __init__(self, hidden_size, mlp_ratio=4.0):
+ super().__init__()
+ self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+ self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+ self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+ self.ca = nn.Sequential(
+ nn.AdaptiveAvgPool2d(1),
+ nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+ nn.Sigmoid(),
+ )
+
+ ffn_channels = int(mlp_ratio * hidden_size)
+ self.conv4 = nn.Conv2d(hidden_size, ffn_channels, 1, bias=True)
+ self.conv5 = nn.Conv2d(ffn_channels, hidden_size, 1, bias=True)
+
+ self.norm1 = MageVAELayerNorm2d(hidden_size, affine=False)
+ self.norm2 = MageVAELayerNorm2d(hidden_size, affine=False)
+
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 6 * hidden_size, bias=True),
+ )
+
+ def forward(self, hidden_states, conditioning):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(conditioning).chunk(
+ 6, dim=1
+ )
+ residual = hidden_states
+ hidden_states = _mage_vae_modulate(self.norm1(residual), shift_msa, scale_msa)
+ hidden_states = F.gelu(self.conv2(self.conv1(hidden_states)))
+ hidden_states = hidden_states * self.ca(hidden_states)
+ hidden_states = self.conv3(hidden_states)
+ hidden_states = residual + gate_msa[..., None, None] * hidden_states
+
+ hidden_states = hidden_states + gate_mlp[..., None, None] * self.conv5(
+ F.gelu(self.conv4(_mage_vae_modulate(self.norm2(hidden_states), shift_mlp, scale_mlp)))
+ )
+ return hidden_states
+
+
+class MageVAEEncoderDiCoBlock(nn.Module):
+ """DConv block without adaLN modulation, for the encoder head pathway."""
+
+ def __init__(self, hidden_size, mlp_ratio=4.0):
+ super().__init__()
+ self.conv1 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+ self.conv2 = nn.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+ self.conv3 = nn.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+ self.ca = nn.Sequential(
+ nn.AdaptiveAvgPool2d(1),
+ nn.Conv2d(hidden_size, hidden_size, 1, bias=True),
+ nn.Sigmoid(),
+ )
+
+ ffn_channels = int(mlp_ratio * hidden_size)
+ self.conv4 = nn.Conv2d(hidden_size, ffn_channels, 1, bias=True)
+ self.conv5 = nn.Conv2d(ffn_channels, hidden_size, 1, bias=True)
+
+ self.norm1 = MageVAELayerNorm2d(hidden_size, affine=True)
+ self.norm2 = MageVAELayerNorm2d(hidden_size, affine=True)
+
+ def forward(self, hidden_states):
+ residual = hidden_states
+ hidden_states = self.norm1(residual)
+ hidden_states = F.gelu(self.conv2(self.conv1(hidden_states)))
+ hidden_states = hidden_states * self.ca(hidden_states)
+ hidden_states = self.conv3(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states + self.conv5(F.gelu(self.conv4(self.norm2(hidden_states))))
+
+
+# ---------------------------------------------------------------------------
+# Nerf-style patch embedder and final layer
+# ---------------------------------------------------------------------------
+class MageVAENerfEmbedder(nn.Module):
+ """Patch-position embedder for the DConv decoder x-pathway."""
+
+ def __init__(self, in_channels, hidden_size_input, max_freqs=8):
+ super().__init__()
+ self.max_freqs = max_freqs
+ self.embedder = nn.Sequential(
+ nn.Linear(in_channels + max_freqs**2, hidden_size_input, bias=True),
+ )
+ self._pos_cache = {}
+
+ def _compute_pos(self, patch_size, device, dtype):
+ key = (patch_size, device, dtype)
+ if key in self._pos_cache:
+ return self._pos_cache[key]
+
+ pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+ pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
+ pos_x = pos_x.reshape(-1, 1, 1)
+ pos_y = pos_y.reshape(-1, 1, 1)
+
+ freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+ fx = freqs[None, :, None]
+ fy = freqs[None, None, :]
+ coeffs = (1 + fx * fy) ** -1
+ dct_x = torch.cos(pos_x * fx * torch.pi)
+ dct_y = torch.cos(pos_y * fy * torch.pi)
+
+ result = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs**2)
+ self._pos_cache[key] = result
+ return result
+
+ def forward(self, x):
+ batch_size, num_patches, _ = x.shape
+ patch_size = int(num_patches**0.5)
+ dct = self._compute_pos(patch_size, x.device, x.dtype).expand(batch_size, -1, -1)
+ return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class MageVAENerfFinalLayer(nn.Module):
+ def __init__(self, hidden_size, out_channels):
+ super().__init__()
+ self.norm = MageVAERMSNorm(hidden_size)
+ self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+ def forward(self, x):
+ return self.linear(self.norm(x))
+
+
+# ---------------------------------------------------------------------------
+# MLP decoder (SimpleMLPAdaLN + MLPResBlock)
+# ---------------------------------------------------------------------------
+class _MageVAEMLPResBlock(nn.Module):
+ def __init__(self, channels):
+ super().__init__()
+ self.in_ln = nn.LayerNorm(channels, eps=1e-6)
+ self.mlp = nn.Sequential(
+ nn.Linear(channels, channels, bias=True),
+ nn.SiLU(),
+ nn.Linear(channels, channels, bias=True),
+ )
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(channels, 3 * channels, bias=True),
+ )
+
+ def forward(self, x, y):
+ shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
+ h = self.in_ln(x) * (1 + scale) + shift
+ return x + gate * self.mlp(h)
+
+
+class MageVAESimpleMLPAdaLN(nn.Module):
+ """Small MLP that maps NerfEmbedder features to per-patch output, conditioned on spatial features."""
+
+ def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size):
+ super().__init__()
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ self.num_res_blocks = num_res_blocks
+ self.patch_size = patch_size
+
+ self.cond_embed = nn.Linear(z_channels, patch_size**2 * model_channels)
+ self.input_proj = nn.Linear(in_channels, model_channels)
+
+ self.res_blocks = nn.ModuleList([_MageVAEMLPResBlock(model_channels) for _ in range(num_res_blocks)])
+
+ def forward(self, x, conditioning):
+ x = self.input_proj(x)
+ conditioning = self.cond_embed(conditioning).reshape(conditioning.shape[0], self.patch_size**2, -1)
+ for block in self.res_blocks:
+ x = block(x, conditioning)
+ return x
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder building blocks (ResNet + Attention)
+# ---------------------------------------------------------------------------
+class MageVAEResnetBlock(nn.Module):
+ """GroupNorm + Conv ResBlock used by the CoD Decoder."""
+
+ def __init__(self, in_channels, out_channels=None, dropout=0.0):
+ super().__init__()
+ out_channels = out_channels or in_channels
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+
+ self.norm1 = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+ self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
+ self.norm2 = nn.GroupNorm(num_groups=32, num_channels=out_channels, eps=1e-6, affine=True)
+ self.dropout = nn.Dropout(dropout)
+ self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
+ if in_channels != out_channels:
+ self.nin_shortcut = nn.Conv2d(in_channels, out_channels, 1)
+
+ def forward(self, x):
+ hidden_states = self.conv1(F.silu(self.norm1(x)))
+ hidden_states = self.conv2(self.dropout(F.silu(self.norm2(hidden_states))))
+ if self.in_channels != self.out_channels:
+ x = self.nin_shortcut(x)
+ return x + hidden_states
+
+
+class MageVAEAttnBlock(nn.Module):
+ """Patched self-attention for the CoD Decoder."""
+
+ def __init__(self, in_channels, patch_size=32):
+ super().__init__()
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+ self.q = nn.Conv2d(in_channels, in_channels, 1)
+ self.k = nn.Conv2d(in_channels, in_channels, 1)
+ self.v = nn.Conv2d(in_channels, in_channels, 1)
+ self.proj_out = nn.Conv2d(in_channels, in_channels, 1)
+
+ def forward(self, x):
+ normalized = self.norm(x)
+ query = self.q(normalized)
+ key = self.k(normalized)
+ value = self.v(normalized)
+
+ d = self.patch_size
+ batch_size, channels, height, width = query.shape
+ pad_h = (d - height % d) % d
+ pad_w = (d - width % d) % d
+ if pad_h or pad_w:
+ query = F.pad(query, (0, pad_w, 0, pad_h), mode="replicate")
+ key = F.pad(key, (0, pad_w, 0, pad_h), mode="replicate")
+ value = F.pad(value, (0, pad_w, 0, pad_h), mode="replicate")
+
+ _, _, height_padded, width_padded = query.shape
+ num_patches_h = height_padded // d
+ num_patches_w = width_padded // d
+ num_patches = num_patches_h * num_patches_w
+
+ # Reshape to patches: [B*num_patches, C, d*d]
+ query = (
+ query.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+ .permute(0, 2, 4, 1, 3, 5)
+ .reshape(batch_size * num_patches, channels, d * d)
+ )
+ key = (
+ key.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+ .permute(0, 2, 4, 1, 3, 5)
+ .reshape(batch_size * num_patches, channels, d * d)
+ )
+ value = (
+ value.reshape(batch_size, channels, num_patches_h, d, num_patches_w, d)
+ .permute(0, 2, 4, 1, 3, 5)
+ .reshape(batch_size * num_patches, channels, d * d)
+ )
+
+ # Attention via F.scaled_dot_product_attention
+ # query/key/value: [B*np, C, d*d] -> [B*np, 1, d*d, C] for SDPA (batch, heads, seq, head_dim)
+ q = query.permute(0, 2, 1).unsqueeze(1)
+ k = key.permute(0, 2, 1).unsqueeze(1)
+ v = value.permute(0, 2, 1).unsqueeze(1)
+ h_ = F.scaled_dot_product_attention(q, k, v, dropout_p=0.0, is_causal=False)
+ h_ = h_.squeeze(1).permute(0, 2, 1) # back to [B*np, C, d*d]
+
+ # Reconstruct
+ hidden_states = (
+ h_
+ .reshape(batch_size, num_patches_h, num_patches_w, channels, d, d)
+ .permute(0, 3, 1, 4, 2, 5)
+ .reshape(batch_size, channels, height_padded, width_padded)
+ )
+ if pad_h or pad_w:
+ hidden_states = hidden_states[:, :, :height, :width]
+
+ return x + self.proj_out(hidden_states)
+
+
+# ---------------------------------------------------------------------------
+# Patch embedding
+# ---------------------------------------------------------------------------
+class MageVAEBottleneckPatchEmbed(nn.Module):
+ """Image patch embed concatenated with a per-patch conditioning vector."""
+
+ def __init__(self, patch_size=16, in_channels=3, bottleneck_dim=128, embed_dim=384, bias=True):
+ super().__init__()
+ self.proj1 = nn.Conv2d(in_channels, bottleneck_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+ self.proj2 = nn.Conv2d(bottleneck_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+ def forward(self, x, conditioning):
+ return self.proj2(torch.cat([self.proj1(x), conditioning], dim=1))
+
+
+# ---------------------------------------------------------------------------
+# adaLN constant folding
+# ---------------------------------------------------------------------------
+class _MageVAEConstAdaLN(nn.Module):
+ """Replaces an adaLN_modulation MLP with a precomputed constant buffer."""
+
+ def __init__(self, modulation: torch.Tensor):
+ super().__init__()
+ self.register_buffer("modulation", modulation.detach().clone())
+
+ def forward(self, conditioning):
+ batch_size = conditioning.shape[0]
+ if self.modulation.shape[0] != batch_size:
+ return self.modulation.expand(batch_size, *self.modulation.shape[1:])
+ return self.modulation
+
+
+# ---------------------------------------------------------------------------
+# DConv Encoder
+# ---------------------------------------------------------------------------
+class MageVAEDConvEncoder(nn.Module):
+ """One-step diffusion encoder: image -> packed (mean, logvar) latent."""
+
+ def __init__(
+ self,
+ latent_channels=128,
+ hidden_size=384,
+ num_blocks=21,
+ patch_size=16,
+ mlp_ratio=4.0,
+ head_size=768,
+ num_head_blocks=2,
+ out_ch_mult=2,
+ ):
+ super().__init__()
+ self.latent_channels = latent_channels
+ self.patch_size = patch_size
+
+ self.patch_cond_embed = nn.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True)
+ self.head_blocks = nn.ModuleList(
+ [MageVAEEncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks)]
+ )
+ self.proj_down = nn.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
+
+ self.z_proj = nn.Conv2d(latent_channels, hidden_size, kernel_size=1, bias=True)
+ self.fuse_proj = nn.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True)
+
+ self.t_embedder = MageVAETimestepEmbedder(hidden_size)
+ self.blocks = nn.ModuleList([MageVAEDiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)])
+
+ self.norm_out = MageVAELayerNorm2d(hidden_size, affine=True)
+ self.proj_out = nn.Conv2d(hidden_size, latent_channels * out_ch_mult, kernel_size=1, bias=True)
+
+ def forward(self, z_t, t, image):
+ conditioning = self.patch_cond_embed(image)
+ for block in self.head_blocks:
+ conditioning = block(conditioning)
+ conditioning = self.proj_down(conditioning)
+
+ hidden_states = self.fuse_proj(torch.cat([conditioning, self.z_proj(z_t)], dim=1))
+ timestep_embedding = self.t_embedder(t.view(-1))
+ for block in self.blocks:
+ hidden_states = block(hidden_states, timestep_embedding)
+ return self.proj_out(self.norm_out(hidden_states))
+
+
+# ---------------------------------------------------------------------------
+# CoD Decoder: latent -> conditioning features for the denoiser
+# ---------------------------------------------------------------------------
+class MageVAEDecoder(nn.Module):
+ """CoD (Cascaded-of-Decoders) decoder: latent -> spatial conditioning features."""
+
+ def __init__(self, out_ch=384, z_ch=128):
+ super().__init__()
+ self.conv_in = nn.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
+ self.block = nn.Sequential(
+ MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+ MageVAEAttnBlock(out_ch, patch_size=32),
+ MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+ MageVAEAttnBlock(out_ch, patch_size=32),
+ MageVAEResnetBlock(in_channels=out_ch, out_channels=out_ch),
+ )
+ self.norm_out = nn.GroupNorm(num_groups=32, num_channels=out_ch, eps=1e-6, affine=True)
+ self.conv_out = nn.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+
+ def forward(self, z):
+ hidden_states = self.block(self.conv_in(z))
+ hidden_states = self.conv_out(F.silu(self.norm_out(hidden_states)))
+ return hidden_states
+
+
+# ---------------------------------------------------------------------------
+# Y-Embedder wrapper (holds the CoD decoder)
+# ---------------------------------------------------------------------------
+class _MageVAEYEmbedder(nn.Module):
+ """Namespace wrapper for the CoD decoder, matching the original checkpoint's
+ ``pipeline.y_embedder.decoder.*`` weight key hierarchy."""
+
+ def __init__(self, hidden_size=384, latent_channels=128):
+ super().__init__()
+ self.decoder = MageVAEDecoder(out_ch=hidden_size, z_ch=latent_channels)
+
+
+# ---------------------------------------------------------------------------
+# DConv Denoiser: conditioning + zero noise -> reconstructed image
+# ---------------------------------------------------------------------------
+class MageVAEDConvDenoiser(nn.Module):
+ """One-step denoiser: takes conditioning from CoD decoder and produces the output image."""
+
+ def __init__(
+ self,
+ patch_size=16,
+ in_channels=3,
+ hidden_size=384,
+ hidden_size_x=32,
+ mlp_ratio=4.0,
+ num_blocks=24,
+ num_cond_blocks=21,
+ bottleneck_dim=128,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.patch_size = patch_size
+ self.hidden_size = hidden_size
+ self.num_cond_blocks = num_cond_blocks
+
+ self.t_embedder = MageVAETimestepEmbedder(hidden_size)
+ self.y_embedder_x = nn.Conv2d(hidden_size, hidden_size_x * patch_size**2, 1, 1, 0)
+ self.x_embedder = MageVAENerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
+ self.s_embedder = MageVAEBottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+
+ self.blocks = nn.ModuleList(
+ [MageVAEDiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)]
+ )
+
+ self.dec_net = MageVAESimpleMLPAdaLN(
+ in_channels=hidden_size_x,
+ model_channels=hidden_size_x,
+ out_channels=in_channels,
+ z_channels=hidden_size,
+ num_res_blocks=num_blocks - num_cond_blocks,
+ patch_size=patch_size,
+ )
+ self.final_layer = MageVAENerfFinalLayer(hidden_size_x, in_channels)
+ self.y_embedder = _MageVAEYEmbedder(hidden_size=hidden_size, latent_channels=bottleneck_dim)
+
+ def forward(self, x, t, conditioning, is_latent=False):
+ if is_latent:
+ conditioning = self.y_embedder.decoder(conditioning)
+
+ batch_size, _, height, width = x.shape
+ timestep_embedding = self.t_embedder(t.view(-1))
+
+ # Spatial conditioning path
+ spatial = self.s_embedder(x, conditioning)
+ for block in self.blocks:
+ spatial = block(spatial, timestep_embedding)
+
+ num_spatial = spatial.shape[-2] * spatial.shape[-1]
+ spatial_flat = spatial.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+ # Per-patch x-pathway
+ x_unfolded = F.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
+ y_x = self.y_embedder_x(conditioning).flatten(2)
+ x_combined = torch.cat([x_unfolded, y_x], dim=1)
+
+ # Reshape: [B, (in_ch + hidden_x), ps^2, num_spatial] -> [B*num_spatial, ps^2, (in_ch + hidden_x)]
+ x_combined = (
+ x_combined.reshape(batch_size, -1, self.patch_size**2, num_spatial).permute(0, 3, 2, 1).flatten(0, 1)
+ )
+
+ x_embedded = self.x_embedder(x_combined)
+ x_decoded = self.dec_net(x_embedded, spatial_flat)
+ x_final = self.final_layer(x_decoded)
+
+ # Fold patches back to image: [B*num_spatial, ps^2, in_ch] -> [B, in_ch, H, W]
+ x_final = x_final.transpose(1, 2).reshape(batch_size, num_spatial, -1)
+ return F.fold(
+ x_final.transpose(1, 2).contiguous(),
+ (height, width),
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ )
+
+
+# ---------------------------------------------------------------------------
+# Main autoencoder
+# ---------------------------------------------------------------------------
+class AutoencoderMageVAE(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ MageVAE autoencoder model using a one-step diffusion encoder and a DConv denoiser
+ with a CoD (Cascaded-of-Decoders) decoder.
+
+ This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods
+ implemented for all models (such as downloading or saving).
+
+ Encoder: DConvEncoder takes an image [B, 3, H, W] and produces a latent [B, 128, H/16, W/16].
+ Decoder: CoD Decoder + DConvDenoiser takes a latent and reconstructs the image [B, 3, H, W].
+
+ Args:
+ latent_channels (`int`, defaults to `128`):
+ Number of channels in the latent space.
+ downsample_factor (`int`, defaults to `16`):
+ Spatial downsampling factor from image to latent.
+ encoder_hidden_size (`int`, defaults to `384`):
+ Hidden dimension of the encoder DConv blocks.
+ encoder_num_blocks (`int`, defaults to `21`):
+ Number of adaLN-modulated DConv blocks in the encoder.
+ encoder_patch_size (`int`, defaults to `16`):
+ Patch size for the encoder's image tokenization.
+ encoder_head_size (`int`, defaults to `768`):
+ Channel dimension of the encoder's head blocks.
+ encoder_num_head_blocks (`int`, defaults to `2`):
+ Number of head blocks in the encoder (without adaLN).
+ decoder_hidden_size (`int`, defaults to `384`):
+ Hidden dimension of the decoder DConv blocks.
+ decoder_hidden_size_x (`int`, defaults to `32`):
+ Hidden dimension of the decoder's per-patch x-pathway.
+ decoder_num_blocks (`int`, defaults to `24`):
+ Total number of blocks in the decoder (cond blocks + MLP res blocks).
+ decoder_num_cond_blocks (`int`, defaults to `21`):
+ Number of adaLN-modulated DConv blocks in the decoder.
+ decoder_bottleneck_dim (`int`, defaults to `128`):
+ Bottleneck dimension for the patch embedding and CoD decoder input.
+ decoder_patch_size (`int`, defaults to `16`):
+ Patch size for the decoder.
+ sample_posterior (`bool`, defaults to `True`):
+ Whether to sample from the posterior (mean + noise * std) or use the mean directly.
+ """
+
+ _no_split_modules = ["MageVAEDiCoBlock", "MageVAEResnetBlock", "MageVAEAttnBlock"]
+ _supports_gradient_checkpointing = False
+
+ @register_to_config
+ def __init__(
+ self,
+ latent_channels: int = 128,
+ downsample_factor: int = 16,
+ encoder_hidden_size: int = 384,
+ encoder_num_blocks: int = 21,
+ encoder_patch_size: int = 16,
+ encoder_head_size: int = 768,
+ encoder_num_head_blocks: int = 2,
+ decoder_hidden_size: int = 384,
+ decoder_hidden_size_x: int = 32,
+ decoder_num_blocks: int = 24,
+ decoder_num_cond_blocks: int = 21,
+ decoder_bottleneck_dim: int = 128,
+ decoder_patch_size: int = 16,
+ sample_posterior: bool = True,
+ ):
+ super().__init__()
+
+ self.encoder = MageVAEDConvEncoder(
+ latent_channels=latent_channels,
+ hidden_size=encoder_hidden_size,
+ num_blocks=encoder_num_blocks,
+ patch_size=encoder_patch_size,
+ head_size=encoder_head_size,
+ num_head_blocks=encoder_num_head_blocks,
+ )
+
+ self.decoder = MageVAEDConvDenoiser(
+ patch_size=decoder_patch_size,
+ in_channels=3,
+ hidden_size=decoder_hidden_size,
+ hidden_size_x=decoder_hidden_size_x,
+ num_blocks=decoder_num_blocks,
+ num_cond_blocks=decoder_num_cond_blocks,
+ bottleneck_dim=decoder_bottleneck_dim,
+ )
+
+ def encode(self, x: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+ """
+ Encode images to latents.
+
+ Args:
+ x (`torch.Tensor`): Input images of shape `[B, 3, H, W]`. H and W must be
+ multiples of `encoder_patch_size`.
+ generator (`torch.Generator`, *optional*):
+ A torch generator for reproducible sampling.
+
+ Returns:
+ `torch.Tensor`: Latent of shape `[B, 128, H/16, W/16]`.
+ """
+ batch_size, _, height, width = x.shape
+ patch_size = self.config.encoder_patch_size
+ latent_channels = self.config.latent_channels
+
+ z_t = torch.zeros(
+ batch_size,
+ latent_channels,
+ height // patch_size,
+ width // patch_size,
+ device=x.device,
+ dtype=x.dtype,
+ )
+ t = torch.zeros(batch_size, device=x.device, dtype=x.dtype)
+
+ out = self.encoder(z_t, t, x)
+ mean = out[:, :latent_channels]
+ logvar = out[:, latent_channels:].clamp(min=-20.0, max=10.0)
+
+ if self.config.sample_posterior:
+ noise = randn_tensor(mean.shape, generator=generator, device=mean.device, dtype=mean.dtype)
+ return mean + torch.exp(0.5 * logvar) * noise
+ return mean
+
+ def forward(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+ return self.decode(z, return_dict=return_dict)
+
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]:
+ """
+ Decode latents to images.
+
+ Args:
+ z (`torch.Tensor`): Latent of shape `[B, 128, H/16, W/16]`.
+ return_dict (`bool`, defaults to `True`):
+ Whether to return a [`~models.autoencoders.vae.DecoderOutput`] or a plain tuple.
+
+ Returns:
+ [`~models.autoencoders.vae.DecoderOutput`] or `tuple`:
+ Decoded images of shape `[B, 3, H, W]`.
+ """
+ batch_size = z.shape[0]
+ height = z.shape[2] * self.config.downsample_factor
+ width = z.shape[3] * self.config.downsample_factor
+ noise = torch.zeros(batch_size, 3, height, width, device=z.device, dtype=z.dtype)
+ t = torch.zeros(batch_size, device=z.device, dtype=z.dtype)
+ sample = self.decoder(noise, t, z, is_latent=True)
+
+ if not return_dict:
+ return (sample,)
+ return DecoderOutput(sample=sample)
+
+ def freeze_adaln(self):
+ """Constant-fold adaLN_modulation MLPs at t=0 for both encoder and decoder.
+
+ At t=0 the adaLN modulation outputs are constant (they only depend on the
+ timestep embedding). This method precomputes those constants and replaces the
+ MLP modules with lightweight buffer wrappers, saving compute and parameters.
+ """
+ device = next(self.parameters()).device
+ dtype = next(self.parameters()).dtype
+ t = torch.zeros(1, device=device, dtype=dtype)
+
+ c_enc = self.encoder.t_embedder(t)
+ count_enc = self._replace_adaln_with_const(self.encoder, c_enc)
+
+ c_dec = self.decoder.t_embedder(t)
+ count_dec = self._replace_adaln_with_const(self.decoder, c_dec)
+
+ logger.info(f"MageVAE: folded {count_enc} encoder + {count_dec} decoder adaLN blocks")
+
+ @staticmethod
+ def _replace_adaln_with_const(module: nn.Module, conditioning: torch.Tensor) -> int:
+ """Replace adaLN_modulation MLPs in MageVAEDiCoBlock instances with constant buffers."""
+ count = 0
+ for child in module.modules():
+ if not isinstance(child, MageVAEDiCoBlock):
+ continue
+ adaln = child.adaLN_modulation
+ if isinstance(adaln, _MageVAEConstAdaLN):
+ continue
+ with torch.no_grad():
+ modulation = adaln(conditioning)
+ child.adaLN_modulation = _MageVAEConstAdaLN(modulation)
+ count += 1
+ return count
diff --git a/pipelines/mageflow/convert_mage_flow_to_diffusers.py b/pipelines/mageflow/convert_mage_flow_to_diffusers.py
new file mode 100644
index 000000000..ad580423a
--- /dev/null
+++ b/pipelines/mageflow/convert_mage_flow_to_diffusers.py
@@ -0,0 +1,269 @@
+"""Convert an original Mage-Flow HF repo layout into diffusers-format weights.
+
+Example
+-------
+
+ python scripts/convert_mage_flow_to_diffusers.py \
+ --input_dir path/to/Mage-Flow-Base \
+ --output_dir path/to/Mage-Flow-Base-diffusers \
+ --dtype bfloat16
+"""
+
+import argparse
+import json
+import os
+import shutil
+from typing import Dict
+
+import safetensors.torch
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Transformer conversion
+# ---------------------------------------------------------------------------
+
+# Top-level 1:1 renames. Anything not matched here (transformer_blocks.*, the
+# time_text_embed / norm_out / proj_out subtrees) is passed through unchanged.
+TRANSFORMER_TOP_LEVEL_RENAMES: Dict[str, str] = {
+ "img_in.weight": "x_embedder.weight",
+ "img_in.bias": "x_embedder.bias",
+ "txt_norm.weight": "context_embedder_norm.weight",
+ "txt_in.weight": "context_embedder.weight",
+ "txt_in.bias": "context_embedder.bias",
+}
+
+
+TRANSFORMER_DIFFUSERS_CONFIG = {
+ "_class_name": "MageFlowTransformer2DModel",
+ "_diffusers_version": "0.37.0",
+ "in_channels": 128,
+ "out_channels": 128,
+ "context_in_dim": 2560,
+ "hidden_size": 3072,
+ "num_attention_heads": 24,
+ "num_layers": 12,
+ "axes_dim": [16, 56, 56],
+ "patch_size": 1,
+}
+
+
+def convert_transformer_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+ new_state_dict: Dict[str, torch.Tensor] = {}
+ for key, tensor in state_dict.items():
+ if key in TRANSFORMER_TOP_LEVEL_RENAMES:
+ new_key = TRANSFORMER_TOP_LEVEL_RENAMES[key]
+ else:
+ # Pass-through: transformer_blocks.*, time_text_embed.*, norm_out.*, proj_out.*
+ new_key = key
+ if new_key in new_state_dict:
+ raise ValueError(f"Duplicate destination key while converting transformer: {new_key}")
+ new_state_dict[new_key] = tensor
+ return new_state_dict
+
+
+# ---------------------------------------------------------------------------
+# VAE conversion
+# ---------------------------------------------------------------------------
+
+VAE_ENCODER_PREFIX = "student.dconv_encoder."
+VAE_DECODER_PREFIX = "pipeline."
+
+# Sub-trees of the original decoder ("pipeline.*") that belong to the Flux2
+# encoder and must be dropped instead of being mapped into the diffusers
+# decoder namespace.
+VAE_DECODER_EXCLUDE_PREFIXES = (
+ "pipeline.y_embedder.encoder.",
+ "pipeline.y_embedder.bottleneck.",
+)
+
+
+VAE_DIFFUSERS_CONFIG = {
+ "_class_name": "AutoencoderMageVAE",
+ "_diffusers_version": "0.37.0",
+ "latent_channels": 128,
+ "downsample_factor": 16,
+ "encoder_hidden_size": 384,
+ "encoder_num_blocks": 21,
+ "encoder_patch_size": 16,
+ "encoder_head_size": 768,
+ "encoder_num_head_blocks": 2,
+ "decoder_hidden_size": 384,
+ "decoder_hidden_size_x": 32,
+ "decoder_num_blocks": 24,
+ "decoder_num_cond_blocks": 21,
+ "decoder_bottleneck_dim": 128,
+ "decoder_patch_size": 16,
+ "sample_posterior": False,
+}
+
+
+def convert_vae_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+ new_state_dict: Dict[str, torch.Tensor] = {}
+ for key, tensor in state_dict.items():
+ if key.startswith(VAE_ENCODER_PREFIX):
+ new_key = "encoder." + key[len(VAE_ENCODER_PREFIX):]
+ elif key.startswith(VAE_DECODER_PREFIX):
+ if any(key.startswith(p) for p in VAE_DECODER_EXCLUDE_PREFIXES):
+ continue
+ # pipeline.y_embedder.decoder.* naturally maps to
+ # decoder.y_embedder.decoder.* under this rule, matching the spec.
+ new_key = "decoder." + key[len(VAE_DECODER_PREFIX):]
+ else:
+ raise ValueError(f"Unexpected VAE key with no known prefix: {key}")
+ if new_key in new_state_dict:
+ raise ValueError(f"Duplicate destination key while converting VAE: {new_key}")
+ new_state_dict[new_key] = tensor
+ return new_state_dict
+
+
+# ---------------------------------------------------------------------------
+# I/O helpers
+# ---------------------------------------------------------------------------
+
+DTYPE_MAP = {
+ "float32": torch.float32,
+ "fp32": torch.float32,
+ "float16": torch.float16,
+ "fp16": torch.float16,
+ "bfloat16": torch.bfloat16,
+ "bf16": torch.bfloat16,
+}
+
+
+def cast_state_dict(state_dict: Dict[str, torch.Tensor], dtype: torch.dtype) -> Dict[str, torch.Tensor]:
+ out: Dict[str, torch.Tensor] = {}
+ for key, tensor in state_dict.items():
+ # Leave integer / bool buffers alone (e.g. num_batches_tracked).
+ if tensor.is_floating_point():
+ out[key] = tensor.to(dtype)
+ else:
+ out[key] = tensor
+ return out
+
+
+def save_component(
+ state_dict: Dict[str, torch.Tensor],
+ config: Dict,
+ output_component_dir: str,
+) -> None:
+ os.makedirs(output_component_dir, exist_ok=True)
+ safetensors.torch.save_file(
+ state_dict,
+ os.path.join(output_component_dir, "diffusion_pytorch_model.safetensors"),
+ )
+ with open(os.path.join(output_component_dir, "config.json"), "w", encoding="utf-8") as f:
+ json.dump(config, f, indent=2)
+ f.write("\n")
+
+
+def copy_scheduler(input_dir: str, output_dir: str) -> None:
+ src = os.path.join(input_dir, "scheduler", "scheduler_config.json")
+ dst_dir = os.path.join(output_dir, "scheduler")
+ os.makedirs(dst_dir, exist_ok=True)
+ shutil.copyfile(src, os.path.join(dst_dir, "scheduler_config.json"))
+
+
+def copy_text_encoder(input_dir: str, output_dir: str, symlink: bool) -> None:
+ src = os.path.join(input_dir, "text_encoder")
+ dst = os.path.join(output_dir, "text_encoder")
+ if os.path.lexists(dst):
+ if os.path.islink(dst) or os.path.isfile(dst):
+ os.remove(dst)
+ else:
+ shutil.rmtree(dst)
+ if symlink:
+ os.symlink(os.path.abspath(src), dst)
+ else:
+ shutil.copytree(src, dst)
+
+
+MODEL_INDEX = {
+ "_class_name": "MageFlowPipeline",
+ "_diffusers_version": "0.37.0",
+ "transformer": ["diffusers", "MageFlowTransformer2DModel"],
+ "vae": ["diffusers", "AutoencoderMageVAE"],
+ "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+ "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"],
+ "tokenizer": ["transformers", "AutoTokenizer"],
+}
+
+
+def write_model_index(output_dir: str) -> None:
+ with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
+ json.dump(MODEL_INDEX, f, indent=2)
+ f.write("\n")
+
+
+# ---------------------------------------------------------------------------
+# Driver
+# ---------------------------------------------------------------------------
+
+
+def parse_args() -> argparse.Namespace:
+ parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+ parser.add_argument("--input_dir", required=True, help="Original Mage-Flow HF repo path.")
+ parser.add_argument("--output_dir", required=True, help="Destination diffusers-format directory.")
+ parser.add_argument("--dtype", default="bfloat16", choices=sorted(DTYPE_MAP.keys()), help="Output tensor dtype.")
+ parser.add_argument(
+ "--text_encoder_mode",
+ default="symlink",
+ choices=["symlink", "copy"],
+ help="How to include the text_encoder directory in the output.",
+ )
+ return parser.parse_args()
+
+
+def main() -> None:
+ args = parse_args()
+ dtype = DTYPE_MAP[args.dtype]
+
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ # --- Transformer ---
+ print("[transformer] loading original weights ...")
+ transformer_sd = safetensors.torch.load_file(
+ os.path.join(args.input_dir, "transformer", "diffusion_pytorch_model.safetensors")
+ )
+ print(f"[transformer] converting {len(transformer_sd)} tensors ...")
+ transformer_sd = convert_transformer_state_dict(transformer_sd)
+ transformer_sd = cast_state_dict(transformer_sd, dtype)
+ save_component(
+ transformer_sd,
+ TRANSFORMER_DIFFUSERS_CONFIG,
+ os.path.join(args.output_dir, "transformer"),
+ )
+ print(f"[transformer] wrote {len(transformer_sd)} tensors to {args.output_dir}/transformer")
+ del transformer_sd
+
+ # --- VAE ---
+ print("[vae] loading original weights ...")
+ vae_sd = safetensors.torch.load_file(
+ os.path.join(args.input_dir, "vae", "diffusion_pytorch_model.safetensors")
+ )
+ print(f"[vae] converting {len(vae_sd)} tensors ...")
+ vae_sd = convert_vae_state_dict(vae_sd)
+ vae_sd = cast_state_dict(vae_sd, dtype)
+ save_component(
+ vae_sd,
+ VAE_DIFFUSERS_CONFIG,
+ os.path.join(args.output_dir, "vae"),
+ )
+ print(f"[vae] wrote {len(vae_sd)} tensors to {args.output_dir}/vae")
+ del vae_sd
+
+ # --- Scheduler ---
+ print("[scheduler] copying config ...")
+ copy_scheduler(args.input_dir, args.output_dir)
+
+ # --- Text encoder ---
+ print(f"[text_encoder] {args.text_encoder_mode} ...")
+ copy_text_encoder(args.input_dir, args.output_dir, symlink=(args.text_encoder_mode == "symlink"))
+
+ # --- model_index.json ---
+ write_model_index(args.output_dir)
+ print(f"Done. Diffusers-format repo written to: {args.output_dir}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pipelines/mageflow/pipeline_mage_flow.py b/pipelines/mageflow/pipeline_mage_flow.py
new file mode 100644
index 000000000..495631635
--- /dev/null
+++ b/pipelines/mageflow/pipeline_mage_flow.py
@@ -0,0 +1,665 @@
+# Copyright 2025 Microsoft and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable
+
+import numpy as np
+import torch
+from transformers import Qwen2Tokenizer, Qwen3VLForConditionalGeneration
+
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from .transformer_mage_flow import MageFlowTransformer2DModel
+from .pipeline_output import MageFlowPipelineOutput
+from .autoencoder_mage_vae import AutoencoderMageVAE
+
+
+if is_torch_xla_available():
+ import torch_xla.core.xla_model as xm
+
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```py
+ >>> import torch
+ >>> from diffusers import MageFlowPipeline
+
+ >>> pipe = MageFlowPipeline.from_pretrained("microsoft/Mage-Flow-4B", torch_dtype=torch.bfloat16)
+ >>> pipe.to("cuda")
+ >>> prompt = "A cat holding a sign that says hello world"
+ >>> image = pipe(prompt, num_inference_steps=30, guidance_scale=5.0).images[0]
+ >>> image.save("mage_flow.png")
+ ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: int | None = None,
+ device: str | torch.device | None = None,
+ timesteps: list[int] | None = None,
+ sigmas: list[float] | None = None,
+ **kwargs,
+):
+ r"""
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`list[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class MageFlowPipeline(DiffusionPipeline):
+ r"""
+ The Mage-Flow pipeline for text-to-image generation.
+
+ Args:
+ transformer ([`MageFlowTransformer2DModel`]):
+ Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+ scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+ A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+ vae ([`AutoencoderMageVAE`]):
+ Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+ text_encoder ([`Qwen3VLForConditionalGeneration`]):
+ Qwen3-VL text encoder for producing text conditioning embeddings.
+ tokenizer (`AutoTokenizer`):
+ Tokenizer for the Qwen3-VL text encoder.
+ """
+
+ model_cpu_offload_seq = "text_encoder->transformer->vae"
+ _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+ def __init__(
+ self,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ vae: AutoencoderMageVAE,
+ text_encoder: Qwen3VLForConditionalGeneration,
+ tokenizer: Qwen2Tokenizer,
+ transformer: MageFlowTransformer2DModel,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ vae=vae,
+ text_encoder=text_encoder,
+ tokenizer=tokenizer,
+ transformer=transformer,
+ scheduler=scheduler,
+ )
+ self.vae_scale_factor = 16 # MageVAE downsample factor
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
+ self.tokenizer_max_length = 2048
+ self.default_sample_size = 64 # 1024 / 16 = 64
+ # ChatML prompt template (same as QwenImage)
+ self.prompt_template = (
+ "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, "
+ "text, spatial relationships of the objects and background:"
+ "<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+ )
+ self.prompt_template_start_idx = 34 # number of system-prompt tokens to skip
+
+ def _get_prompt_embeds(
+ self,
+ prompt: str | list[str] | None = None,
+ device: torch.device | None = None,
+ dtype: torch.dtype | None = None,
+ ):
+ device = device or self._execution_device
+ if self.text_encoder is None:
+ raise ValueError(
+ "Text encoder is not available. Please provide `prompt_embeds` directly "
+ "when the pipeline is initialized without a text encoder."
+ )
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+
+ template = self.prompt_template
+ drop_idx = self.prompt_template_start_idx
+ txt = [template.format(e) for e in prompt]
+ txt_tokens = self.tokenizer(
+ txt, max_length=self.tokenizer_max_length + drop_idx, padding=True, truncation=True, return_tensors="pt"
+ ).to(device)
+ encoder_out = self.text_encoder(
+ input_ids=txt_tokens.input_ids,
+ attention_mask=txt_tokens.attention_mask,
+ output_hidden_states=True,
+ )
+ hidden_states = encoder_out.hidden_states[-1]
+
+ # Extract valid tokens per sample, drop system prompt prefix, then re-pad to uniform length.
+ bool_mask = txt_tokens.attention_mask.bool()
+ valid_lengths = bool_mask.sum(dim=1)
+ selected = hidden_states[bool_mask]
+ split_hidden_states = torch.split(selected, valid_lengths.tolist(), dim=0)
+ split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+
+ attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+ max_seq_len = max([e.size(0) for e in split_hidden_states])
+ prompt_embeds = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+ )
+ prompt_embeds_mask = torch.stack(
+ [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+ )
+
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ return prompt_embeds, prompt_embeds_mask
+
+ def encode_prompt(
+ self,
+ prompt: str | list[str],
+ device: torch.device | None = None,
+ num_images_per_prompt: int = 1,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ max_sequence_length: int = 2048,
+ ):
+ r"""
+ Encode the text prompt into embeddings for the transformer.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ Prompt to be encoded.
+ device (`torch.device`):
+ Torch device.
+ num_images_per_prompt (`int`):
+ Number of images that should be generated per prompt.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ prompt_embeds_mask (`torch.Tensor`, *optional*):
+ Attention mask for `prompt_embeds`.
+ max_sequence_length (`int`):
+ Maximum sequence length for the text embeddings.
+ """
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds, prompt_embeds_mask = self._get_prompt_embeds(prompt, device)
+
+ prompt_embeds = prompt_embeds[:, :max_sequence_length]
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+
+ if prompt_embeds_mask is not None:
+ prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length]
+ prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+ if prompt_embeds_mask.all():
+ prompt_embeds_mask = None
+
+ return prompt_embeds, prompt_embeds_mask
+
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt=None,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ prompt_embeds_mask=None,
+ negative_prompt_embeds_mask=None,
+ callback_on_step_end_tensor_inputs=None,
+ max_sequence_length=None,
+ ):
+ if height % self.vae_scale_factor != 0 or width % self.vae_scale_factor != 0:
+ logger.warning(
+ f"`height` and `width` have to be divisible by {self.vae_scale_factor} but are {height} and {width}. "
+ "Dimensions will be resized accordingly"
+ )
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found "
+ f"{[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and prompt_embeds_mask is None:
+ logger.warning(
+ "`prompt_embeds` is provided and `prompt_embeds_mask` is not provided, so the model will treat all"
+ " prompt tokens as valid. If `prompt_embeds` contains padding, you should provide the padding mask as"
+ " `prompt_embeds_mask`. Make sure to generate `prompt_embeds_mask` from the same text encoder that was"
+ " used to generate `prompt_embeds`."
+ )
+
+ if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+ logger.warning(
+ "`negative_prompt_embeds` is provided and `negative_prompt_embeds_mask` is not provided, so the model"
+ " will treat all negative prompt tokens as valid. If `negative_prompt_embeds` contains padding, you"
+ " should provide the padding mask as `negative_prompt_embeds_mask`. Make sure to generate"
+ " `negative_prompt_embeds_mask` from the same text encoder that was used to generate"
+ " `negative_prompt_embeds`."
+ )
+
+ if max_sequence_length is not None and max_sequence_length > 2048:
+ raise ValueError(f"`max_sequence_length` cannot be greater than 2048 but is {max_sequence_length}")
+
+ @staticmethod
+ def _prepare_latent_image_ids(height, width, device, dtype):
+ latent_image_ids = torch.zeros(height, width, 3, device=device, dtype=dtype)
+ latent_image_ids[..., 1] = torch.arange(height, device=device, dtype=dtype)[:, None]
+ latent_image_ids[..., 2] = torch.arange(width, device=device, dtype=dtype)[None, :]
+ latent_image_ids = latent_image_ids.reshape(height * width, 3)
+ return latent_image_ids
+
+ def prepare_latents(
+ self,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ # MageVAE: 16x downsample, no patch packing
+ height = height // self.vae_scale_factor
+ width = width // self.vae_scale_factor
+
+ shape = (batch_size, num_channels_latents, height, width)
+
+ if latents is not None:
+ return latents.to(device=device, dtype=dtype)
+
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ # Flatten to sequence: [B, C, H, W] -> [B, H*W, C]
+ latents = latents.permute(0, 2, 3, 1).reshape(batch_size, height * width, num_channels_latents)
+
+ return latents
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def attention_kwargs(self):
+ return self._attention_kwargs
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def current_timestep(self):
+ return self._current_timestep
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: str | list[str] | None = None,
+ negative_prompt: str | list[str] | None = None,
+ height: int | None = None,
+ width: int | None = None,
+ num_inference_steps: int = 30,
+ guidance_scale: float = 5.0,
+ num_images_per_prompt: int = 1,
+ generator: torch.Generator | list[torch.Generator] | None = None,
+ latents: torch.Tensor | None = None,
+ prompt_embeds: torch.Tensor | None = None,
+ prompt_embeds_mask: torch.Tensor | None = None,
+ negative_prompt_embeds: torch.Tensor | None = None,
+ negative_prompt_embeds_mask: torch.Tensor | None = None,
+ output_type: str | None = "pil",
+ return_dict: bool = True,
+ attention_kwargs: dict[str, Any] | None = None,
+ callback_on_step_end: Callable[[int, int], None] | None = None,
+ callback_on_step_end_tensor_inputs: list[str] = ["latents"],
+ max_sequence_length: int = 2048,
+ sigmas: list[float] | None = None,
+ ) -> MageFlowPipelineOutput | tuple:
+ r"""
+ Function invoked when calling the pipeline for generation.
+
+ Args:
+ prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+ instead.
+ negative_prompt (`str` or `list[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ not greater than `1`).
+ height (`int`, *optional*, defaults to `self.default_sample_size * self.vae_scale_factor`):
+ The height in pixels of the generated image.
+ width (`int`, *optional*, defaults to `self.default_sample_size * self.vae_scale_factor`):
+ The width in pixels of the generated image.
+ num_inference_steps (`int`, *optional*, defaults to 30):
+ The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+ expense of slower inference.
+ guidance_scale (`float`, *optional*, defaults to 5.0):
+ Classifier-free guidance scale. Enabled by setting `guidance_scale > 1`. Higher guidance scale
+ encourages images closely linked to the text `prompt`, usually at the expense of lower image quality.
+ num_images_per_prompt (`int`, *optional*, defaults to 1):
+ The number of images to generate per prompt.
+ generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+ One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+ to make generation deterministic.
+ latents (`torch.Tensor`, *optional*):
+ Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+ generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+ tensor will be generated by sampling using the supplied random `generator`.
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ prompt_embeds_mask (`torch.Tensor`, *optional*):
+ Attention mask for `prompt_embeds`.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ negative_prompt_embeds_mask (`torch.Tensor`, *optional*):
+ Attention mask for `negative_prompt_embeds`.
+ output_type (`str`, *optional*, defaults to `"pil"`):
+ The output format of the generated image. Choose between
+ [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+ return_dict (`bool`, *optional*, defaults to `True`):
+ Whether or not to return a [`~pipelines.mage_flow.MageFlowPipelineOutput`] instead of a plain tuple.
+ attention_kwargs (`dict`, *optional*):
+ A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+ `self.processor` in
+ [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+ callback_on_step_end (`Callable`, *optional*):
+ A function that calls at the end of each denoising steps during the inference. The function is called
+ with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+ callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+ `callback_on_step_end_tensor_inputs`.
+ callback_on_step_end_tensor_inputs (`list`, *optional*):
+ The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+ will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+ `._callback_tensor_inputs` attribute of your pipeline class.
+ max_sequence_length (`int`, defaults to 2048):
+ Maximum sequence length to use with the `prompt`.
+ sigmas (`list[float]`, *optional*):
+ Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+ their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+ will be used.
+
+ Examples:
+
+ Returns:
+ [`~pipelines.mage_flow.MageFlowPipelineOutput`] or `tuple`:
+ [`~pipelines.mage_flow.MageFlowPipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
+ returning a tuple, the first element is a list with the generated images.
+ """
+
+ height = height or self.default_sample_size * self.vae_scale_factor
+ width = width or self.default_sample_size * self.vae_scale_factor
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt=negative_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+ callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+ max_sequence_length=max_sequence_length,
+ )
+
+ self._guidance_scale = guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._current_timestep = None
+ self._interrupt = False
+
+ # 2. Define call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # 3. Encode prompt
+ prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+ prompt=prompt,
+ prompt_embeds=prompt_embeds,
+ prompt_embeds_mask=prompt_embeds_mask,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ )
+ if do_classifier_free_guidance:
+ if negative_prompt_embeds is None and self.text_encoder is None:
+ # text_encoder unavailable and no negative_prompt_embeds provided, skip CFG
+ do_classifier_free_guidance = False
+ else:
+ negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+ prompt=negative_prompt if negative_prompt is not None else [""] * batch_size,
+ prompt_embeds=negative_prompt_embeds,
+ prompt_embeds_mask=negative_prompt_embeds_mask,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ max_sequence_length=max_sequence_length,
+ )
+
+ # 4. Prepare latent variables
+ num_channels_latents = self.transformer.config.in_channels
+ latents = self.prepare_latents(
+ batch_size * num_images_per_prompt,
+ num_channels_latents,
+ height,
+ width,
+ prompt_embeds.dtype,
+ device,
+ generator,
+ latents,
+ )
+
+ # 5. Prepare image position ids for RoPE
+ latent_h = height // self.vae_scale_factor
+ latent_w = width // self.vae_scale_factor
+ img_ids = self._prepare_latent_image_ids(latent_h, latent_w, device, prompt_embeds.dtype)
+
+ # 6. Prepare timesteps
+ # Mage-Flow uses base_sigmas = linspace(1, 1/N, N) which differs from the scheduler's
+ # default sigma computation. The scheduler's shift=6.0 is applied on top of these.
+ if sigmas is None:
+ sigmas = np.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps)
+ timesteps, num_inference_steps = retrieve_timesteps(
+ self.scheduler,
+ num_inference_steps,
+ device,
+ sigmas=sigmas,
+ )
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ self._num_timesteps = len(timesteps)
+
+ if self.attention_kwargs is None:
+ self._attention_kwargs = {}
+
+ # 7. Denoising loop
+ self.scheduler.set_begin_index(0)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ self._current_timestep = t
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+ if do_classifier_free_guidance:
+ noise_pred_cond = self.transformer(
+ hidden_states=latents,
+ encoder_hidden_states=prompt_embeds,
+ timestep=timestep / 1000,
+ img_ids=img_ids,
+ joint_attention_kwargs=self.attention_kwargs,
+ return_dict=False,
+ )[0]
+ noise_pred_uncond = self.transformer(
+ hidden_states=latents,
+ encoder_hidden_states=negative_prompt_embeds,
+ timestep=timestep / 1000,
+ img_ids=img_ids,
+ joint_attention_kwargs=self.attention_kwargs,
+ return_dict=False,
+ )[0]
+ noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond)
+ else:
+ noise_pred = self.transformer(
+ hidden_states=latents,
+ encoder_hidden_states=prompt_embeds,
+ timestep=timestep / 1000,
+ img_ids=img_ids,
+ joint_attention_kwargs=self.attention_kwargs,
+ return_dict=False,
+ )[0]
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents_dtype = latents.dtype
+ latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+ if latents.dtype != latents_dtype:
+ if torch.backends.mps.is_available():
+ # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+ latents = latents.to(latents_dtype)
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+ if XLA_AVAILABLE:
+ xm.mark_step()
+
+ self._current_timestep = None
+
+ # 8. VAE decode
+ if output_type == "latent":
+ image = latents
+ else:
+ # Unflatten: [B, H*W, C] -> [B, C, H, W]
+ latents = latents.reshape(batch_size * num_images_per_prompt, latent_h, latent_w, num_channels_latents)
+ latents = latents.permute(0, 3, 1, 2)
+ latents = latents.to(self.vae.dtype)
+ image = self.vae(latents, return_dict=False)[0]
+ image = image.clamp(-1, 1)
+ image = self.image_processor.postprocess(image, output_type=output_type)
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ return (image,)
+
+ return MageFlowPipelineOutput(images=image)
diff --git a/pipelines/mageflow/pipeline_output.py b/pipelines/mageflow/pipeline_output.py
new file mode 100644
index 000000000..87b58bd09
--- /dev/null
+++ b/pipelines/mageflow/pipeline_output.py
@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+
+import numpy as np
+import PIL.Image
+
+from diffusers.utils import BaseOutput
+
+
+@dataclass
+class MageFlowPipelineOutput(BaseOutput):
+ """
+ Output class for Mage-Flow pipelines.
+
+ Args:
+ images (`list[PIL.Image.Image]` or `np.ndarray`)
+ List of denoised PIL images of length `batch_size` or numpy array of shape `(batch_size, height, width,
+ num_channels)`. PIL images or numpy array present the denoised images of the diffusion pipeline.
+ """
+
+ images: list[PIL.Image.Image] | np.ndarray
diff --git a/pipelines/mageflow/transformer_mage_flow.py b/pipelines/mageflow/transformer_mage_flow.py
new file mode 100644
index 000000000..9464a4231
--- /dev/null
+++ b/pipelines/mageflow/transformer_mage_flow.py
@@ -0,0 +1,630 @@
+# Copyright 2025 The Mage Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import logging
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin, FeedForward
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import TimestepEmbedding
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+def _apply_rotary_emb_complex(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor:
+ """Apply complex rotary embeddings to ``x`` using MageFlow's adjacent-pair convention.
+
+ Args:
+ x: Query or key tensor of shape ``[B, S, H, D]``.
+ freqs_cis: Complex frequency tensor of shape ``[S, D_rope // 2]`` where
+ ``D_rope = sum(axes_dim)``. When ``D_rope < D`` only the first
+ ``D_rope`` dimensions are rotated; the rest pass through unchanged.
+
+ Returns:
+ Tensor of same shape and dtype as *x* with rotary embeddings applied.
+ """
+ rope_dim = freqs_cis.shape[-1] * 2 # complex dim -> real dim
+ head_dim = x.shape[-1]
+
+ if rope_dim < head_dim:
+ x_rope = x[..., :rope_dim]
+ x_pass = x[..., rope_dim:]
+ else:
+ x_rope = x
+ x_pass = None
+
+ # [B, S, H, rope_dim] -> [B, S, H, rope_dim/2] complex
+ x_complex = torch.view_as_complex(x_rope.float().reshape(*x_rope.shape[:-1], -1, 2))
+ # freqs_cis: [S, D_rope/2] -> [1, S, 1, D_rope/2] for broadcasting
+ freqs = freqs_cis.unsqueeze(0).unsqueeze(2)
+ x_rotated = torch.view_as_real(x_complex * freqs).flatten(-2)
+ x_rotated = x_rotated.to(x.dtype)
+
+ if x_pass is not None:
+ return torch.cat([x_rotated, x_pass], dim=-1)
+ return x_rotated
+
+
+class MageFlowPosEmbed(nn.Module):
+ """Complex RoPE with symmetric positive/negative frequency scaling for MageFlow.
+
+ Computes multi-scale rotary positional embeddings for video/image tokens using
+ three axes (frame, height, width). Height and width axes use symmetric
+ positive/negative frequency indices centered around the spatial midpoint.
+ """
+
+ def __init__(self, theta: int = 10000, axes_dim: list[int] | None = None):
+ super().__init__()
+ if axes_dim is None:
+ axes_dim = [16, 48, 48]
+ self.theta = theta
+ self.axes_dim = axes_dim
+
+ pos_index = torch.arange(4096)
+ neg_index = torch.arange(4096).flip(0) * -1 - 1
+
+ pos_freqs = torch.cat(
+ [
+ self._rope_params(pos_index, self.axes_dim[0], self.theta),
+ self._rope_params(pos_index, self.axes_dim[1], self.theta),
+ self._rope_params(pos_index, self.axes_dim[2], self.theta),
+ ],
+ dim=1,
+ )
+ neg_freqs = torch.cat(
+ [
+ self._rope_params(neg_index, self.axes_dim[0], self.theta),
+ self._rope_params(neg_index, self.axes_dim[1], self.theta),
+ self._rope_params(neg_index, self.axes_dim[2], self.theta),
+ ],
+ dim=1,
+ )
+ self.register_buffer("pos_freqs_real", pos_freqs.real.contiguous(), persistent=False)
+ self.register_buffer("pos_freqs_imag", pos_freqs.imag.contiguous(), persistent=False)
+ self.register_buffer("neg_freqs_real", neg_freqs.real.contiguous(), persistent=False)
+ self.register_buffer("neg_freqs_imag", neg_freqs.imag.contiguous(), persistent=False)
+
+ @staticmethod
+ def _rope_params(index: torch.Tensor, dim: int, theta: float = 10000.0) -> torch.Tensor:
+ """Compute complex RoPE frequencies for a 1-D position index."""
+ freqs = torch.outer(
+ index.float(),
+ 1.0 / torch.pow(theta, torch.arange(0, dim, 2, dtype=torch.float32).div(dim)),
+ )
+ return torch.polar(torch.ones_like(freqs), freqs)
+
+ def _compute_video_freqs(self, frame: int, height: int, width: int, idx: int = 0) -> torch.Tensor:
+ seq_len = frame * height * width
+ pos_freqs = torch.complex(self.pos_freqs_real.float(), self.pos_freqs_imag.float())
+ neg_freqs = torch.complex(self.neg_freqs_real.float(), self.neg_freqs_imag.float())
+ freqs_pos = pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+ freqs_neg = neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+ freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+ freqs_height = torch.cat(
+ [freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]],
+ dim=0,
+ )
+ freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+ freqs_width = torch.cat(
+ [freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]],
+ dim=0,
+ )
+ freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+
+ freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_len, -1)
+ return freqs.clone().contiguous()
+
+ def forward(self, img_ids: torch.Tensor, height: int | None = None, width: int | None = None) -> torch.Tensor:
+ """Compute RoPE frequencies from image position ids.
+
+ Args:
+ img_ids: ``[seq_len, 3]`` tensor with (frame, height, width) position
+ indices for each image token.
+ height: Latent spatial height.
+ width: Latent spatial width.
+
+ Returns:
+ Complex frequency tensor of shape ``[seq_len, head_dim // 2]``.
+ """
+ frame = 1
+ freqs = self._compute_video_freqs(frame, height, width, idx=0)
+ return freqs.to(img_ids.device)
+
+ @staticmethod
+ def _infer_grid_size(img_ids: torch.Tensor) -> tuple[int, int]:
+ height = int(img_ids[:, 1].max().item()) + 1
+ width = int(img_ids[:, 2].max().item()) + 1
+ return height, width
+
+
+class MageFlowTimestepProjEmbeddings(nn.Module):
+ """Timestep projection embeddings for MageFlow.
+
+ Uses a custom sinusoidal embedding that downcasts the frequency table to the
+ input dtype before computing the embedding. The model was trained with this
+ exact bf16 rounding, so using diffusers' standard float32 variant degrades
+ output quality.
+ """
+
+ def __init__(self, embedding_dim: int):
+ super().__init__()
+ self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+ self.num_channels = 256
+ self.scale = 1000
+
+ @staticmethod
+ def _sinusoidal_embedding(
+ timesteps: torch.Tensor,
+ embedding_dim: int,
+ scale: float = 1.0,
+ max_period: int = 10000,
+ ) -> torch.Tensor:
+ half_dim = embedding_dim // 2
+ exponent = -math.log(max_period) * torch.arange(
+ start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+ )
+ exponent = exponent / half_dim
+
+ # Downcast frequency table to input dtype (bf16) before multiplying —
+ # the model was trained with this exact rounding.
+ emb = torch.exp(exponent).to(timesteps.dtype)
+ emb = timesteps[:, None].float() * emb[None, :]
+ emb = scale * emb
+ emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+ # flip sin to cos
+ emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+ return emb
+
+ def forward(self, timestep: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor:
+ timesteps_proj = self._sinusoidal_embedding(timestep, self.num_channels, scale=self.scale)
+ timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))
+ return timesteps_emb
+
+
+class MageFlowAttnProcessor:
+ """Attention processor for MageFlow double-stream (MMDiT) architecture.
+
+ Implements joint attention over concatenated ``[text, image]`` tokens. RoPE is
+ applied only to image query/key, not text.
+ """
+
+ _attention_backend = None
+ _parallel_config = None
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(f"{self.__class__.__name__} requires PyTorch 2.0. Please upgrade your pytorch version.")
+
+ def __call__(
+ self,
+ attn: "MageFlowAttention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ # Compute QKV for image stream
+ img_query = attn.to_q(hidden_states)
+ img_key = attn.to_k(hidden_states)
+ img_value = attn.to_v(hidden_states)
+
+ # Reshape to multi-head: [B, S, inner_dim] -> [B, S, H, D]
+ img_query = img_query.unflatten(-1, (attn.heads, -1))
+ img_key = img_key.unflatten(-1, (attn.heads, -1))
+ img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+ # Apply QK normalization
+ img_query = attn.norm_q(img_query)
+ img_key = attn.norm_k(img_key)
+
+ # Apply RoPE to image Q/K only (not text)
+ if image_rotary_emb is not None:
+ img_query = _apply_rotary_emb_complex(img_query, image_rotary_emb)
+ img_key = _apply_rotary_emb_complex(img_key, image_rotary_emb)
+
+ if encoder_hidden_states is not None and attn.added_kv_proj_dim is not None:
+ # Compute QKV for text stream
+ txt_query = attn.add_q_proj(encoder_hidden_states)
+ txt_key = attn.add_k_proj(encoder_hidden_states)
+ txt_value = attn.add_v_proj(encoder_hidden_states)
+
+ txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+ txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+ txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+ txt_query = attn.norm_added_q(txt_query)
+ txt_key = attn.norm_added_k(txt_key)
+
+ # No RoPE on text — concatenate [text, image] for joint attention
+ query = torch.cat([txt_query, img_query], dim=1)
+ key = torch.cat([txt_key, img_key], dim=1)
+ value = torch.cat([txt_value, img_value], dim=1)
+ else:
+ query = img_query
+ key = img_key
+ value = img_value
+
+ # Joint attention via dispatch
+ attn_output = dispatch_attention_fn(
+ query,
+ key,
+ value,
+ attn_mask=attention_mask,
+ backend=self._attention_backend,
+ parallel_config=self._parallel_config,
+ )
+ attn_output = attn_output.flatten(2, 3)
+ attn_output = attn_output.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ # Split back into text and image parts
+ txt_seq_len = encoder_hidden_states.shape[1]
+ txt_attn_output, img_attn_output = attn_output.split_with_sizes(
+ [txt_seq_len, attn_output.shape[1] - txt_seq_len], dim=1
+ )
+ img_attn_output = attn.to_out[0](img_attn_output)
+ img_attn_output = attn.to_out[1](img_attn_output)
+ txt_attn_output = attn.to_add_out(txt_attn_output)
+ return img_attn_output, txt_attn_output
+
+ return attn_output
+
+
+class MageFlowAttention(nn.Module, AttentionModuleMixin):
+ """Multi-head attention module for MageFlow with support for dual-stream (MMDiT) attention.
+
+ Follows the diffusers attention pattern with ``_default_processor_cls`` and
+ ``_available_processors`` for backend dispatch.
+ """
+
+ _default_processor_cls = MageFlowAttnProcessor
+ _available_processors = [MageFlowAttnProcessor]
+
+ def __init__(
+ self,
+ query_dim: int,
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = True,
+ added_kv_proj_dim: int | None = None,
+ added_proj_bias: bool | None = True,
+ out_bias: bool = True,
+ eps: float = 1e-6,
+ out_dim: int | None = None,
+ elementwise_affine: bool = True,
+ processor: "MageFlowAttnProcessor | None" = None,
+ ):
+ super().__init__()
+ self.head_dim = dim_head
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.dropout = dropout
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.added_proj_bias = added_proj_bias
+
+ self.norm_q = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.RMSNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(nn.Dropout(dropout))
+
+ if added_kv_proj_dim is not None:
+ self.norm_added_q = nn.RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = nn.RMSNorm(dim_head, eps=eps)
+ self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.to_add_out = nn.Linear(self.inner_dim, query_dim, bias=out_bias)
+
+ if processor is None:
+ processor = self._default_processor_cls()
+ self.set_processor(processor)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ image_rotary_emb: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ return self.processor(self, hidden_states, encoder_hidden_states, attention_mask, image_rotary_emb, **kwargs)
+
+
+class MageFlowTransformerBlock(nn.Module):
+ """Double-stream MMDiT transformer block for MageFlow.
+
+ Each block processes image and text streams with separate modulation (AdaLN),
+ joint attention, and separate feed-forward networks.
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ num_attention_heads: int,
+ attention_head_dim: int,
+ eps: float = 1e-6,
+ ):
+ super().__init__()
+ self.dim = dim
+ self.num_attention_heads = num_attention_heads
+ self.attention_head_dim = attention_head_dim
+
+ # Image stream modulation and layers
+ self.img_mod = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(dim, 6 * dim, bias=True),
+ )
+ self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.attn = MageFlowAttention(
+ query_dim=dim,
+ added_kv_proj_dim=dim,
+ dim_head=attention_head_dim,
+ heads=num_attention_heads,
+ out_dim=dim,
+ bias=True,
+ processor=MageFlowAttnProcessor(),
+ eps=eps,
+ )
+ self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ # Text stream modulation and layers
+ self.txt_mod = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(dim, 6 * dim, bias=True),
+ )
+ self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+ self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ temb: torch.Tensor,
+ image_rotary_emb: torch.Tensor | None = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # Compute modulation parameters for both streams
+ img_mod_params = self.img_mod(temb)
+ txt_mod_params = self.txt_mod(temb)
+
+ # Split into norm1 and norm2 modulation parameters (each has shift, scale, gate)
+ img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)
+ txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)
+
+ # Image stream: norm1 + modulation
+ img_shift1, img_scale1, img_gate1 = img_mod1.chunk(3, dim=-1)
+ img_normed = self.img_norm1(hidden_states)
+ img_modulated = img_normed * (1 + img_scale1.unsqueeze(1)) + img_shift1.unsqueeze(1)
+
+ # Text stream: norm1 + modulation
+ txt_shift1, txt_scale1, txt_gate1 = txt_mod1.chunk(3, dim=-1)
+ txt_normed = self.txt_norm1(encoder_hidden_states)
+ txt_modulated = txt_normed * (1 + txt_scale1.unsqueeze(1)) + txt_shift1.unsqueeze(1)
+
+ # Joint attention
+ joint_attention_kwargs = joint_attention_kwargs or {}
+ img_attn_output, txt_attn_output = self.attn(
+ hidden_states=img_modulated,
+ encoder_hidden_states=txt_modulated,
+ image_rotary_emb=image_rotary_emb,
+ **joint_attention_kwargs,
+ )
+
+ # Apply gates and residuals
+ hidden_states = hidden_states + img_gate1.unsqueeze(1) * img_attn_output
+ encoder_hidden_states = encoder_hidden_states + txt_gate1.unsqueeze(1) * txt_attn_output
+
+ # Image stream: norm2 + MLP
+ img_shift2, img_scale2, img_gate2 = img_mod2.chunk(3, dim=-1)
+ img_normed2 = self.img_norm2(hidden_states)
+ img_modulated2 = img_normed2 * (1 + img_scale2.unsqueeze(1)) + img_shift2.unsqueeze(1)
+ img_mlp_output = self.img_mlp(img_modulated2)
+ hidden_states = hidden_states + img_gate2.unsqueeze(1) * img_mlp_output
+
+ # Text stream: norm2 + MLP
+ txt_shift2, txt_scale2, txt_gate2 = txt_mod2.chunk(3, dim=-1)
+ txt_normed2 = self.txt_norm2(encoder_hidden_states)
+ txt_modulated2 = txt_normed2 * (1 + txt_scale2.unsqueeze(1)) + txt_shift2.unsqueeze(1)
+ txt_mlp_output = self.txt_mlp(txt_modulated2)
+ encoder_hidden_states = encoder_hidden_states + txt_gate2.unsqueeze(1) * txt_mlp_output
+
+ # Clip to prevent overflow for fp16
+ if encoder_hidden_states.dtype == torch.float16:
+ encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+ if hidden_states.dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return encoder_hidden_states, hidden_states
+
+
+class MageFlowTransformer2DModel(
+ ModelMixin,
+ ConfigMixin,
+ PeftAdapterMixin,
+ FromOriginalModelMixin,
+ CacheMixin,
+ AttentionMixin,
+):
+ """Transformer model for MageFlow image generation.
+
+ A dual-stream (MMDiT) Transformer that processes image and text tokens jointly.
+ Uses complex multi-scale RoPE for image positional encoding and Qwen3-VL text
+ embeddings as conditioning.
+
+ Args:
+ in_channels (`int`, defaults to ``128``):
+ Number of channels in the input latent (MageVAE latent channels).
+ out_channels (`int`, defaults to ``128``):
+ Number of channels in the output.
+ context_in_dim (`int`, defaults to ``3584``):
+ Dimension of the text encoder hidden states (Qwen3-VL hidden size).
+ hidden_size (`int`, defaults to ``3072``):
+ Inner dimension of the transformer (num_attention_heads * attention_head_dim).
+ num_attention_heads (`int`, defaults to ``24``):
+ Number of attention heads.
+ num_layers (`int`, defaults to ``32``):
+ Number of dual-stream transformer blocks.
+ axes_dim (`list[int]``, defaults to ``[16, 48, 48]``):
+ RoPE dimension split across axes (frame, height, width). Must sum to
+ ``hidden_size // num_attention_heads``.
+ patch_size (`int`, defaults to ``1``):
+ Patch size for the output projection.
+ """
+
+ _supports_gradient_checkpointing = True
+ _no_split_modules = ["MageFlowTransformerBlock"]
+ _repeated_blocks = ["MageFlowTransformerBlock"]
+ _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+ main_input_name = "hidden_states"
+
+ @register_to_config
+ def __init__(
+ self,
+ in_channels: int = 128,
+ out_channels: int = 128,
+ context_in_dim: int = 3584,
+ hidden_size: int = 3072,
+ num_attention_heads: int = 24,
+ num_layers: int = 32,
+ axes_dim: list[int] = [16, 48, 48],
+ patch_size: int = 1,
+ ):
+ super().__init__()
+ self.out_channels = out_channels
+ self.inner_dim = hidden_size
+ self.num_attention_heads = num_attention_heads
+ attention_head_dim = hidden_size // num_attention_heads
+
+ self.pos_embed = MageFlowPosEmbed(theta=10000, axes_dim=axes_dim)
+
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim)
+ self.context_embedder_norm = nn.RMSNorm(context_in_dim, eps=1e-6)
+ self.context_embedder = nn.Linear(context_in_dim, self.inner_dim)
+
+ self.time_text_embed = MageFlowTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ MageFlowTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+
+ self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor = None,
+ timestep: torch.Tensor = None,
+ img_ids: torch.Tensor = None,
+ joint_attention_kwargs: dict[str, Any] | None = None,
+ return_dict: bool = True,
+ ) -> torch.Tensor | Transformer2DModelOutput:
+ """
+ The [`MageFlowTransformer2DModel`] forward method.
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch_size, img_seq_len, in_channels)`):
+ Flattened image latent tokens.
+ encoder_hidden_states (`torch.Tensor` of shape `(batch_size, txt_seq_len, context_in_dim)`):
+ Text encoder hidden states (Qwen3-VL embeddings).
+ timestep (`torch.Tensor`):
+ Raw sigma value in ``[0, 1]``.
+ img_ids (`torch.Tensor` of shape `(img_seq_len, 3)`):
+ Image position ids ``(frame, height, width)`` for RoPE computation.
+ joint_attention_kwargs (`dict`, *optional*):
+ Additional keyword arguments passed to the attention processor.
+ return_dict (`bool`, defaults to ``True``):
+ Whether to return a :class:`Transformer2DModelOutput` or a plain tuple.
+
+ Returns:
+ :class:`Transformer2DModelOutput` or ``tuple``.
+ """
+ # Embed image tokens
+ hidden_states = self.x_embedder(hidden_states)
+
+ # Embed text tokens: RMSNorm then linear projection
+ encoder_hidden_states = self.context_embedder_norm(encoder_hidden_states)
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+ # Timestep embedding (Timesteps module handles the 1000x scaling internally via scale=1000)
+ timestep = timestep.to(hidden_states.dtype)
+ temb = self.time_text_embed(timestep, hidden_states)
+
+ # Compute image RoPE (text tokens are not rotated)
+ if img_ids.ndim == 3:
+ img_ids = img_ids[0]
+ image_rotary_emb = self.pos_embed(img_ids, *MageFlowPosEmbed._infer_grid_size(img_ids))
+
+ # Transformer blocks
+ for block in self.transformer_blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+ encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+ block,
+ hidden_states,
+ encoder_hidden_states,
+ temb,
+ image_rotary_emb,
+ joint_attention_kwargs,
+ )
+ else:
+ encoder_hidden_states, hidden_states = block(
+ hidden_states=hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ temb=temb,
+ image_rotary_emb=image_rotary_emb,
+ joint_attention_kwargs=joint_attention_kwargs,
+ )
+
+ # Final norm and projection (image stream only)
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
+
+ if not return_dict:
+ return (output,)
+
+ return Transformer2DModelOutput(sample=output)
diff --git a/pipelines/model_anima.py b/pipelines/model_anima.py
index 488b0074a..cbf10885c 100644
--- a/pipelines/model_anima.py
+++ b/pipelines/model_anima.py
@@ -58,48 +58,6 @@ def load_anima(checkpoint_info, diffusers_load_config=None):
if repo_id is None or repo_id.lower() == 'none':
return None
- # load-or-download custom pipeline modules from repo
- """
- import os
- import sys
- import huggingface_hub as hf
-
- if os.path.exists(os.path.join(repo_id, 'pipeline.py')):
- pipeline_file = os.path.join(repo_id, 'pipeline.py')
- else:
- try:
- if os.path.exists(repo_id):
- from pipelines.generic_map import transformers_map
- custom_id = transformers_map.get('AnimaTextToImagePipeline', repo_id)
- else:
- custom_id = repo_id
- pipeline_file = hf.hf_hub_download(repo_id=custom_id, filename='pipeline.py', cache_dir=shared.opts.hfcache_dir)
- except Exception as e:
- log.error(f'Load model: type=Anima failed to download custom modules: {e}')
- return None
- if os.path.exists(os.path.join(repo_id, 'llm_adapter/modeling_llm_adapter.py')):
- adapter_file = os.path.join(repo_id, 'llm_adapter/modeling_llm_adapter.py')
- else:
- try:
- if os.path.exists(repo_id):
- from pipelines.generic_map import transformers_map
- custom_id = transformers_map.get('AnimaTextToImagePipeline', repo_id)
- else:
- custom_id = repo_id
- adapter_file = hf.hf_hub_download(repo_id=custom_id, filename='llm_adapter/modeling_llm_adapter.py', cache_dir=shared.opts.hfcache_dir)
- except Exception as e:
- log.error(f'Load model: type=Anima failed to download custom modules: {e}')
- return None
-
- # dynamically import custom classes and register in sys.modules so Diffusers' from_pretrained can resolve them via trust_remote_code
- adapter_mod = _import_from_file('modeling_llm_adapter', adapter_file)
- sys.modules['modeling_llm_adapter'] = adapter_mod
- pipeline_mod = _import_from_file('pipeline', pipeline_file)
- sys.modules['pipeline'] = pipeline_mod
- AnimaTextToImagePipeline = pipeline_mod.AnimaTextToImagePipeline
- AnimaLLMAdapter = adapter_mod.AnimaLLMAdapter
- """
-
import sys
from pipelines.anima import modeling_llm_adapter
sys.modules['modeling_llm_adapter'] = modeling_llm_adapter
diff --git a/pipelines/model_boogu.py b/pipelines/model_boogu.py
index fa62ac75c..e6c9dfd34 100644
--- a/pipelines/model_boogu.py
+++ b/pipelines/model_boogu.py
@@ -22,20 +22,20 @@ def load_boogu(checkpoint_info, diffusers_load_config=None):
from pipelines.boogu import transformer_boogu, scheduling_flow_match_euler_discrete_time_shifting
sys.modules['transformer_boogu'] = transformer_boogu # for loading custom code from HF repo
sys.modules['scheduling_flow_match_euler_discrete_time_shifting'] = scheduling_flow_match_euler_discrete_time_shifting # for loading custom code from HF repo
- scheduler = scheduling_flow_match_euler_discrete_time_shifting.FlowMatchEulerDiscreteScheduler.from_pretrained(repo_id, subfolder='scheduler', cache_dir=shared.opts.diffusers_dir)
+ generic.set_pipeline('Boogu', BooguImagePipeline)
if repo_id is None or repo_id.lower() == 'none':
return None
mllm = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLForConditionalGeneration, load_config=diffusers_load_config, subfolder='mllm')
transformer = generic.load_transformer(repo_id, cls_name=BooguImageTransformer2DModel, load_config=diffusers_load_config)
+ scheduler = scheduling_flow_match_euler_discrete_time_shifting.FlowMatchEulerDiscreteScheduler.from_pretrained(repo_id, subfolder='scheduler', cache_dir=shared.opts.diffusers_dir)
if 'turbo' in repo_id.lower():
cls = BooguImageTurboPipeline
else:
cls = BooguImagePipeline
- generic.set_pipeline('Boogu', cls)
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['boogu'] = cls
diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['boogu'] = cls
diff --git a/pipelines/model_ernie.py b/pipelines/model_ernie.py
index 5f9e5a5a8..b739b7fce 100644
--- a/pipelines/model_ernie.py
+++ b/pipelines/model_ernie.py
@@ -14,6 +14,10 @@ def load_ernie_image(checkpoint_info, diffusers_load_config=None):
load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
log.debug(f'Load model: type=ERNIE-Image repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args} pe={shared.opts.model_ernie_enable_pe}')
+ if 'nunchaku-lite' in repo_id.lower():
+ from modules.attention import hijack_kernels
+ hijack_kernels()
+
from pipelines.ernie import ERNIE_SPEC
transformer = generic.load_transformer(
repo_id,
diff --git a/pipelines/model_flux.py b/pipelines/model_flux.py
index cc823163a..fcf4cc44b 100644
--- a/pipelines/model_flux.py
+++ b/pipelines/model_flux.py
@@ -46,6 +46,10 @@ def load_flux(checkpoint_info, diffusers_load_config=None):
from pipelines.flux.flux_nunchaku import load_flux_nunchaku
transformer = load_flux_nunchaku(repo_id)
+ if 'nunchaku-lite' in repo_id.lower():
+ from modules.attention import hijack_kernels
+ hijack_kernels()
+
# finally load transformer and text encoder if not already loaded
if transformer is None:
transformer = generic.load_transformer(repo_id, cls_name=diffusers.FluxTransformer2DModel, load_config=diffusers_load_config)
diff --git a/pipelines/model_krea2.py b/pipelines/model_krea2.py
index afe3ec5ac..4c65fb2e1 100644
--- a/pipelines/model_krea2.py
+++ b/pipelines/model_krea2.py
@@ -15,16 +15,19 @@ def load_krea2(checkpoint_info, diffusers_load_config=None):
from pipelines.krea2.transformer_krea2 import Krea2Transformer2DModel
from pipelines.krea2.pipeline_krea2 import Krea2Pipeline, Krea2Img2ImgPipeline
+ from pipelines.krea2.pipeline_krea2_inpaint import Krea2InpaintPipeline
from pipelines.krea2 import KREA2_SPEC
diffusers.Krea2Transformer2DModel = Krea2Transformer2DModel
diffusers.Krea2Pipeline = Krea2Pipeline
diffusers.Krea2Img2ImgPipeline = Krea2Img2ImgPipeline
+ diffusers.Krea2InpaintPipeline = Krea2InpaintPipeline
generic.set_pipeline('Krea2', Krea2Pipeline)
# One class per task so get_diffusers_task defaults to text2image and set_diffuser_pipe switches
# to the img2img variant cleanly (matches the Chroma/Qwen per-task-class pattern).
from diffusers.pipelines import auto_pipeline
auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Pipeline
auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['krea2'] = Krea2Img2ImgPipeline
+ auto_pipeline.AUTO_INPAINT_PIPELINES_MAPPING['krea2'] = Krea2InpaintPipeline
if repo_id is None or repo_id.lower() == 'none':
return None
diff --git a/pipelines/model_mageflow.py b/pipelines/model_mageflow.py
new file mode 100644
index 000000000..e9ec72fe8
--- /dev/null
+++ b/pipelines/model_mageflow.py
@@ -0,0 +1,55 @@
+import diffusers
+import transformers
+from modules import shared, devices, sd_models, model_quant, sd_hijack_te, sd_hijack_vae
+from modules.logger import log
+from pipelines import generic
+
+
+def load_mageflow(checkpoint_info, diffusers_load_config=None):
+ if diffusers_load_config is None:
+ diffusers_load_config = {}
+
+ repo_id = sd_models.path_to_repo(checkpoint_info)
+ sd_models.hf_auth_check(checkpoint_info)
+
+ load_args, _ = model_quant.get_dit_args(diffusers_load_config, allow_quant=False)
+ log.debug(f'Load model: type=MageFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+
+ from pipelines.mageflow import MageFlowPipeline, MageFlowTransformer2DModel
+
+ log.debug(f'Load model: type=MageFlow repo="{repo_id}" config={diffusers_load_config} offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={diffusers_load_config}')
+ generic.set_pipeline('MageFlow', MageFlowPipeline)
+
+ if repo_id is None or repo_id.lower() == 'none':
+ return None
+
+ transformer = generic.load_transformer(repo_id, cls_name=MageFlowTransformer2DModel, load_config=diffusers_load_config)
+ text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLForConditionalGeneration, load_config=diffusers_load_config)
+ tokenizer = transformers.Qwen2Tokenizer.from_pretrained(repo_id, subfolder="text_encoder", cache_dir=shared.opts.diffusers_dir)
+
+ diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING['mageflow'] = MageFlowPipeline
+ diffusers.pipelines.auto_pipeline.AUTO_IMAGE2IMAGE_PIPELINES_MAPPING['mageflow'] = MageFlowPipeline
+
+ pipe = MageFlowPipeline.from_pretrained(
+ repo_id,
+ transformer=transformer,
+ text_encoder=text_encoder,
+ tokenizer=tokenizer,
+ cache_dir=shared.opts.diffusers_dir,
+ **load_args,
+ )
+
+ pipe.task_args = {
+ 'output_type': 'pil',
+ }
+
+ generic.load_vae_override(pipe, diffusers_load_config)
+ del transformer
+ del text_encoder
+ del tokenizer
+
+ sd_hijack_te.init_hijack(pipe)
+ sd_hijack_vae.init_hijack(pipe)
+
+ devices.torch_gc(force=True, reason='load')
+ return pipe
diff --git a/pipelines/model_prx.py b/pipelines/model_prx.py
index 1fbaba01a..5a35e601d 100644
--- a/pipelines/model_prx.py
+++ b/pipelines/model_prx.py
@@ -20,6 +20,7 @@ def load_prx(checkpoint_info, diffusers_load_config=None):
if repo_id is None or repo_id.lower() == 'none':
return None
+
pipe = diffusers.PRXPipeline.from_pretrained(
repo_id,
transformer=transformer,
diff --git a/pipelines/model_qwen.py b/pipelines/model_qwen.py
index 5dbdf8736..0e44c76d7 100644
--- a/pipelines/model_qwen.py
+++ b/pipelines/model_qwen.py
@@ -16,6 +16,10 @@ def load_qwen(checkpoint_info, diffusers_load_config=None):
load_args, _quant_args = model_quant.get_dit_args(diffusers_load_config, module='Model')
log.debug(f'Load model: type=Qwen model="{checkpoint_info.name}" repo="{repo_id}" offload={shared.opts.diffusers_offload_mode} dtype={devices.dtype} args={load_args}')
+ if 'nunchaku-lite' in repo_id.lower():
+ from modules.attention import hijack_kernels
+ hijack_kernels()
+
if '2509' in repo_id or '2511' in repo_id:
cls_name = diffusers.QwenImageEditPlusPipeline
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["qwen-image"] = diffusers.QwenImageEditPlusPipeline
diff --git a/pipelines/model_sefi.py b/pipelines/model_sefi.py
index aa62b46db..e1cf8f237 100644
--- a/pipelines/model_sefi.py
+++ b/pipelines/model_sefi.py
@@ -18,6 +18,7 @@ def load_sefi(checkpoint_info, diffusers_load_config=None):
transformer = generic.load_transformer(repo_id, cls_name=SeFiTransformer2DModel, load_config=diffusers_load_config)
text_encoder = generic.load_text_encoder(repo_id, cls_name=transformers.Qwen3VLForConditionalGeneration, load_config=diffusers_load_config)
+ generic.set_pipeline('SeFi', SeFiPipeline)
if repo_id is None or repo_id.lower() == 'none':
return None
@@ -31,6 +32,10 @@ def load_sefi(checkpoint_info, diffusers_load_config=None):
diffusers.pipelines.auto_pipeline.AUTO_TEXT2IMAGE_PIPELINES_MAPPING["sefi"] = SeFiPipeline
+ pipe.task_args = {
+ "output_type": "np",
+ }
+
generic.load_vae_override(pipe, diffusers_load_config)
del text_encoder
diff --git a/pipelines/model_z_image.py b/pipelines/model_z_image.py
index dab801f42..c5ada8492 100644
--- a/pipelines/model_z_image.py
+++ b/pipelines/model_z_image.py
@@ -43,6 +43,9 @@ def load_z_image(checkpoint_info, diffusers_load_config=None):
transformer = None
if model_quant.check_nunchaku('Model'): # only available model
transformer = init_nunchaku()
+ if 'nunchaku-lite' in repo_id.lower():
+ from modules.attention import hijack_kernels
+ hijack_kernels()
if transformer is None:
transformer = generic.load_transformer(repo_id, cls_name=diffusers.ZImageTransformer2DModel, load_config=diffusers_load_config)
diff --git a/pipelines/sefi/__init__.py b/pipelines/sefi/__init__.py
index 282265fa0..efd4d82d2 100644
--- a/pipelines/sefi/__init__.py
+++ b/pipelines/sefi/__init__.py
@@ -1,3 +1,7 @@
from .transformer_sefi import SeFiTransformer2DModel
from .pipeline_sefi import SeFiPipeline
from .pipeline_output import SeFiPipelineOutput
+
+import diffusers
+diffusers.SeFiTransformer2DModel = SeFiTransformer2DModel
+diffusers.SeFiPipeline = SeFiPipeline
diff --git a/pipelines/sefi/convert_sefi_to_diffusers.py b/pipelines/sefi/convert_sefi_to_diffusers.py
new file mode 100644
index 000000000..fca38a5a5
--- /dev/null
+++ b/pipelines/sefi/convert_sefi_to_diffusers.py
@@ -0,0 +1,424 @@
+#!/usr/bin/env python
+# Copyright 2026 SeFi-Image Authors and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+
+import torch
+import yaml
+from huggingface_hub import snapshot_download
+from safetensors.torch import load_file
+
+from diffusers import FlowMatchEulerDiscreteScheduler, __version__
+from transformer_sefi import SeFiTransformer2DModel
+
+
+SEFI_SCALE_PRESETS = {
+ "0p5b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 12,
+ "num_layers": 3,
+ "num_single_layers": 10,
+ "joint_attention_dim": 6144,
+ },
+ "1b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 16,
+ "num_layers": 4,
+ "num_single_layers": 12,
+ "joint_attention_dim": 6144,
+ },
+ "2b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 20,
+ "num_layers": 4,
+ "num_single_layers": 16,
+ "joint_attention_dim": 6144,
+ },
+ "3b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 22,
+ "num_layers": 5,
+ "num_single_layers": 18,
+ "joint_attention_dim": 7680,
+ },
+ "4b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 24,
+ "num_layers": 5,
+ "num_single_layers": 20,
+ "joint_attention_dim": 7680,
+ },
+ "5b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 26,
+ "num_layers": 6,
+ "num_single_layers": 21,
+ "joint_attention_dim": 7680,
+ },
+ "6b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 28,
+ "num_layers": 6,
+ "num_single_layers": 22,
+ "joint_attention_dim": 7680,
+ },
+ "8b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 30,
+ "num_layers": 7,
+ "num_single_layers": 24,
+ "joint_attention_dim": 7680,
+ },
+ "9b": {
+ "attention_head_dim": 128,
+ "num_attention_heads": 32,
+ "num_layers": 8,
+ "num_single_layers": 24,
+ "joint_attention_dim": 12288,
+ },
+}
+
+QWEN3VL_TEXT_HIDDEN_DIMS = {
+ "qwen3vl_2b": 2048,
+ "qwen3vl_4b": 2560,
+ "qwen3vl_8b": 4096,
+}
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Convert a SeFi-Image checkpoint to Diffusers format.")
+ parser.add_argument("--checkpoint", required=True, help="Local checkpoint folder or Hugging Face repo id.")
+ parser.add_argument("--output", required=True, help="Output Diffusers checkpoint folder.")
+ parser.add_argument("--cache-dir", default=None, help="Optional Hugging Face cache directory.")
+ parser.add_argument("--token", default=None, help="Optional Hugging Face token for gated checkpoints.")
+ parser.add_argument(
+ "--source-repo-id",
+ default=None,
+ help="Original Hub repo id to record in the converted model card when converting a local checkpoint.",
+ )
+ parser.add_argument(
+ "--target-repo-id",
+ default=None,
+ help="Converted Hub repo id to use in the generated model card example.",
+ )
+ parser.add_argument(
+ "--variant",
+ choices=["base", "rl", "turbo"],
+ default=None,
+ help="Model family. Inferred from checkpoint name if omitted.",
+ )
+ return parser.parse_args()
+
+
+def resolve_checkpoint(checkpoint: str, cache_dir: str | None, token: str | None) -> Path:
+ path = Path(checkpoint).expanduser()
+ if path.exists():
+ return path
+ return Path(snapshot_download(checkpoint, cache_dir=cache_dir, token=token))
+
+
+def copytree(src: Path, dst: Path, ignore=None):
+ if dst.exists():
+ shutil.rmtree(dst)
+ shutil.copytree(src, dst, ignore=ignore)
+
+
+def load_json(path: Path):
+ with open(path, "r", encoding="utf-8") as handle:
+ return json.load(handle)
+
+
+def load_yaml(path: Path):
+ with open(path, "r", encoding="utf-8") as handle:
+ return yaml.safe_load(handle)
+
+
+def save_json(path: Path, payload: dict):
+ with open(path, "w", encoding="utf-8") as handle:
+ json.dump(payload, handle, indent=2, sort_keys=True)
+ handle.write("\n")
+
+
+def save_model_card(output: Path, source_repo_id: str | None, target_repo_id: str, variant: str):
+ metadata = [
+ "---",
+ "license: cc-by-nc-4.0",
+ "library_name: diffusers",
+ "pipeline_tag: text-to-image",
+ "gated: true",
+ "tags:",
+ "- sefi-image",
+ "- semantic-first-diffusion",
+ "- safetensors",
+ ]
+ if source_repo_id is not None:
+ metadata.append(f"base_model: {source_repo_id}")
+ metadata.append("---")
+
+ source_link = (
+ f"[`{source_repo_id}`](https://huggingface.co/{source_repo_id})"
+ if source_repo_id is not None
+ else "the original SeFi-Image checkpoint"
+ )
+ inference_call = (
+ """image = pipe(
+ \"A red apple on a wooden table.\",
+ num_inference_steps=4,
+ guidance_scale=1.0,
+).images[0]"""
+ if variant == "turbo"
+ else 'image = pipe("A red apple on a wooden table.").images[0]'
+ )
+ card = (
+ "\n".join(metadata)
+ + f"""
+
+# SeFi-Image Diffusers checkpoint
+
+This repository is a Diffusers-format conversion of {source_link}. The original checkpoint is not modified by the
+conversion. Refer to the source model card for model details, limitations, and responsible-use guidance.
+
+```python
+import torch
+from diffusers import SeFiPipeline
+
+pipe = SeFiPipeline.from_pretrained(
+ \"{target_repo_id}\", dtype=torch.bfloat16
+).to(\"cuda\")
+{inference_call}
+image.save(\"sefi.png\")
+```
+
+## License
+
+The checkpoint is distributed under the Creative Commons Attribution-NonCommercial 4.0 International license
+(CC BY-NC 4.0). It is for non-commercial use only.
+"""
+ )
+ with open(output / "README.md", "w", encoding="utf-8") as handle:
+ handle.write(card)
+
+
+def infer_variant(checkpoint: str, config: dict, explicit_variant: str | None) -> str:
+ if explicit_variant is not None:
+ return explicit_variant
+ configured = str(config.get("inference", {}).get("family", "") or config.get("model", {}).get("variant", ""))
+ text = f"{checkpoint} {configured}".lower()
+ if "turbo" in text or "distill" in text:
+ return "turbo"
+ if "rl" in text:
+ return "rl"
+ return "base"
+
+
+def default_steps(variant: str) -> int:
+ return 4 if variant == "turbo" else 50
+
+
+def default_guidance_scale(variant: str) -> float:
+ return 1.0 if variant == "turbo" else 4.0
+
+
+def texture_vae_config_path(root: Path, texture_vae_name: str) -> Path:
+ if texture_vae_name in {"sd1.5", "flux1", "flux2"}:
+ return root / "vae" / "config.json"
+ raise ValueError(f"Unsupported texture VAE: {texture_vae_name}")
+
+
+def build_transformer_config(root: Path, sefi_config: dict) -> dict:
+ model_config = sefi_config["model"]
+ transformer_config = load_json(root / "transformer" / "config.json")
+ transformer_config.pop("_class_name", None)
+ transformer_config.pop("_diffusers_version", None)
+ transformer_config.pop("_name_or_path", None)
+ transformer_config.pop("guidance_embeds", None)
+
+ scale = str(model_config.get("transformer_scale", "")).lower()
+ if scale and scale != "custom":
+ transformer_config.update(SEFI_SCALE_PRESETS[scale])
+ elif scale == "custom":
+ transformer_config.update(model_config.get("transformer_overrides", {}))
+
+ semantic_channels = int(model_config["semantic_channels"])
+ texture_vae_name = str(model_config["texture_vae"]["name"]).lower()
+ vae_config = load_json(texture_vae_config_path(root, texture_vae_name))
+ texture_channels = int(vae_config["latent_channels"]) * 4
+ total_channels = semantic_channels + texture_channels
+
+ text_config = model_config["text_encoder"]
+ hidden_layers = tuple(int(layer) for layer in text_config["hidden_layers"])
+ text_dim = int(QWEN3VL_TEXT_HIDDEN_DIMS[text_config["model_name"]]) * len(hidden_layers)
+
+ transformer_config["in_channels"] = total_channels
+ transformer_config["out_channels"] = total_channels
+ if int(transformer_config["joint_attention_dim"]) != text_dim:
+ raise ValueError(
+ "Text dimension mismatch: "
+ f"transformer joint_attention_dim={transformer_config['joint_attention_dim']} vs text_dim={text_dim}."
+ )
+ return transformer_config
+
+
+def load_transformer_state_dict(transformer_dir: Path) -> dict[str, torch.Tensor]:
+ index_path = transformer_dir / "diffusion_pytorch_model.safetensors.index.json"
+ single_path = transformer_dir / "diffusion_pytorch_model.safetensors"
+ bin_path = transformer_dir / "diffusion_pytorch_model.bin"
+
+ if index_path.exists():
+ index = load_json(index_path)
+ state_dict = {}
+ for shard in sorted(set(index["weight_map"].values())):
+ state_dict.update(load_file(transformer_dir / shard))
+ return state_dict
+ if single_path.exists():
+ return load_file(single_path)
+ if bin_path.exists():
+ return torch.load(bin_path, map_location="cpu")
+ raise FileNotFoundError(f"No supported transformer weights found under {transformer_dir}.")
+
+
+def convert_transformer_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+ converted = {}
+ for key, value in state_dict.items():
+ if key.startswith("backbone."):
+ converted_key = key.removeprefix("backbone.")
+ elif key.startswith("dual_time_embed."):
+ converted_key = key
+ else:
+ raise ValueError(f"Unexpected transformer key in the original SeFi checkpoint: {key}")
+
+ if converted_key in converted:
+ raise ValueError(f"Transformer key collision after conversion: {converted_key}")
+ converted[converted_key] = value
+ return converted
+
+
+def convert_scheduler(root: Path, output: Path):
+ scheduler_config = load_json(root / "scheduler" / "scheduler_config.json")
+ source_scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+ scheduler_config["shift"] = 1.0
+ scheduler_config["use_dynamic_shifting"] = False
+ scheduler = FlowMatchEulerDiscreteScheduler.from_config(scheduler_config)
+
+ torch.testing.assert_close(scheduler.sigmas, source_scheduler.sigmas, rtol=0.0, atol=0.0)
+ torch.testing.assert_close(scheduler.timesteps, source_scheduler.timesteps, rtol=0.0, atol=0.0)
+ num_train_timesteps = int(scheduler.config.num_train_timesteps)
+ expected_sigmas = torch.linspace(1.0, 1.0 / num_train_timesteps, num_train_timesteps)
+ torch.testing.assert_close(scheduler.sigmas, expected_sigmas, rtol=0.0, atol=1e-7)
+ torch.testing.assert_close(scheduler.timesteps, expected_sigmas * num_train_timesteps, rtol=0.0, atol=1e-4)
+ scheduler.save_pretrained(output / "scheduler")
+
+
+def copy_tokenizer_files(src: Path, dst: Path):
+ weight_patterns = {
+ "model*.safetensors",
+ "pytorch_model*.bin",
+ "*.index.json",
+ }
+
+ def ignore(_dir, names):
+ ignored = set()
+ for name in names:
+ for pattern in weight_patterns:
+ if Path(name).match(pattern):
+ ignored.add(name)
+ return ignored
+
+ copytree(src, dst, ignore=ignore)
+
+
+def main():
+ args = parse_args()
+ root = resolve_checkpoint(args.checkpoint, args.cache_dir, args.token)
+ output = Path(args.output).expanduser()
+ output.mkdir(parents=True, exist_ok=True)
+
+ sefi_config = load_yaml(root / "sefi_config.yaml")
+ variant = infer_variant(args.checkpoint, sefi_config, args.variant)
+ transformer_config = build_transformer_config(root, sefi_config)
+
+ state_dict = convert_transformer_state_dict(load_transformer_state_dict(root / "transformer"))
+ with torch.device("meta"):
+ transformer = SeFiTransformer2DModel(**transformer_config)
+ transformer.load_state_dict(state_dict, strict=True, assign=True)
+ expected_transformer_keys = set(state_dict)
+ expected_transformer_dtypes = {key: value.dtype for key, value in state_dict.items()}
+ floating_dtypes = {value.dtype for value in state_dict.values() if value.is_floating_point()}
+ if len(floating_dtypes) != 1:
+ raise ValueError(f"Expected one floating-point transformer dtype, got {sorted(map(str, floating_dtypes))}.")
+ transformer_dtype = floating_dtypes.pop()
+ transformer.save_pretrained(output / "transformer", safe_serialization=True)
+ del transformer, state_dict
+
+ reloaded_transformer = SeFiTransformer2DModel.from_pretrained(output / "transformer", dtype=transformer_dtype)
+ if set(reloaded_transformer.state_dict()) != expected_transformer_keys:
+ raise ValueError("Transformer state dict keys changed after the save/load round trip.")
+ round_trip_dtypes = {key: value.dtype for key, value in reloaded_transformer.state_dict().items()}
+ if round_trip_dtypes != expected_transformer_dtypes:
+ raise ValueError("Transformer state dict dtypes changed after the save/load round trip.")
+ del reloaded_transformer
+
+ convert_scheduler(root, output)
+ copytree(root / "vae", output / "vae")
+
+ text_encoder_name = sefi_config["model"]["text_encoder"]["model_name"]
+ qwen_dir_name = {
+ "qwen3vl_2b": "Qwen3-VL-2B-Instruct",
+ "qwen3vl_4b": "Qwen3-VL-4B-Instruct",
+ "qwen3vl_8b": "Qwen3-VL-8B-Instruct",
+ }[text_encoder_name]
+ qwen_dir = root / qwen_dir_name
+ copytree(qwen_dir, output / "text_encoder")
+ copy_tokenizer_files(qwen_dir, output / "tokenizer")
+
+ model_config = sefi_config["model"]
+ inference_config = sefi_config.get("inference", {})
+ training_sefi_config = sefi_config.get("training", {}).get("sefi", {})
+ texture_vae_name = str(model_config["texture_vae"]["name"]).lower()
+ vae_class = "AutoencoderKLFlux2" if texture_vae_name == "flux2" else "AutoencoderKL"
+ model_index = {
+ "_class_name": "SeFiPipeline",
+ "_diffusers_version": __version__,
+ "transformer": ["diffusers", "SeFiTransformer2DModel"],
+ "scheduler": ["diffusers", "FlowMatchEulerDiscreteScheduler"],
+ "vae": ["diffusers", vae_class],
+ "text_encoder": ["transformers", "Qwen3VLForConditionalGeneration"],
+ "tokenizer": ["transformers", "Qwen2Tokenizer"],
+ "semantic_channels": int(model_config["semantic_channels"]),
+ "texture_vae_name": texture_vae_name,
+ "is_turbo": variant == "turbo",
+ "default_guidance_scale": float(inference_config.get("guidance_scale", default_guidance_scale(variant))),
+ "default_num_inference_steps": int(inference_config.get("steps", default_steps(variant))),
+ "delta_t": float(inference_config.get("delta_t", training_sefi_config.get("delta_t_max", 0.1))),
+ "timestep_shift_alpha": float(
+ inference_config.get("timestep_shift_alpha", 1.0 if variant == "turbo" else 0.3)
+ ),
+ "text_encoder_hidden_layers": [int(layer) for layer in model_config["text_encoder"]["hidden_layers"]],
+ "max_sequence_length": int(model_config["text_encoder"].get("max_length", 1024)),
+ }
+ save_json(output / "model_index.json", model_index)
+ source_repo_id = args.source_repo_id
+ if source_repo_id is None and not Path(args.checkpoint).expanduser().exists():
+ source_repo_id = args.checkpoint
+ save_model_card(output, source_repo_id, args.target_repo_id or output.name, variant)
+ shutil.copy2(root / "sefi_config.yaml", output / "sefi_config.yaml")
+ print(f"Saved SeFi-Image Diffusers checkpoint to {output}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/pipelines/sefi/pipeline_sefi.py b/pipelines/sefi/pipeline_sefi.py
index 520d5bb34..1711981ce 100644
--- a/pipelines/sefi/pipeline_sefi.py
+++ b/pipelines/sefi/pipeline_sefi.py
@@ -19,7 +19,7 @@ from transformers import Qwen2Tokenizer, Qwen3VLForConditionalGeneration
from diffusers.models import AutoencoderKL, AutoencoderKLFlux2
from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils import is_torch_xla_available, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
@@ -35,16 +35,15 @@ else:
XLA_AVAILABLE = False
-logger = logging.get_logger(__name__) # pylint: disable=invalid-name
-
-
EXAMPLE_DOC_STRING = """
Examples:
```py
>>> import torch
>>> from diffusers import SeFiPipeline
- >>> pipe = SeFiPipeline.from_pretrained("./sefi-1b-base-diffusers", torch_dtype=torch.bfloat16)
+ >>> pipe = SeFiPipeline.from_pretrained(
+ ... "SeFi-Image/SeFi-Image-1B-Base-diffusers", dtype=torch.bfloat16
+ ... )
>>> pipe.to("cuda")
>>> image = pipe("A red apple on a wooden table.").images[0]
>>> image.save("sefi.png")
@@ -135,13 +134,8 @@ class SeFiPipeline(DiffusionPipeline):
)
if isinstance(text_encoder_hidden_layers, str):
text_encoder_hidden_layers = tuple(int(layer) for layer in text_encoder_hidden_layers.split(","))
- semantic_channels = 16 if semantic_channels is None else semantic_channels
- if texture_vae_name is None:
- texture_vae_name = "flux2" if vae is not None and hasattr(vae, "bn") else "sd1.5"
- default_guidance_scale = 4.0 if default_guidance_scale is None else default_guidance_scale
- default_num_inference_steps = 50 if default_num_inference_steps is None else default_num_inference_steps
- text_encoder_hidden_layers = (9, 18, 27) if text_encoder_hidden_layers is None else text_encoder_hidden_layers
- max_sequence_length = 1024 if max_sequence_length is None else max_sequence_length
+ elif text_encoder_hidden_layers is not None:
+ text_encoder_hidden_layers = tuple(text_encoder_hidden_layers)
self.register_to_config(
semantic_channels=semantic_channels,
texture_vae_name=texture_vae_name,
@@ -150,11 +144,11 @@ class SeFiPipeline(DiffusionPipeline):
default_num_inference_steps=default_num_inference_steps,
delta_t=delta_t,
timestep_shift_alpha=timestep_shift_alpha,
- text_encoder_hidden_layers=tuple(text_encoder_hidden_layers),
+ text_encoder_hidden_layers=text_encoder_hidden_layers,
max_sequence_length=max_sequence_length,
)
- self.semantic_channels = int(semantic_channels)
+ self.semantic_channels = semantic_channels
self.texture_vae_name = str(texture_vae_name).lower()
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8
self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
@@ -184,24 +178,23 @@ class SeFiPipeline(DiffusionPipeline):
def num_timesteps(self):
return self._num_timesteps
- @staticmethod
- def _prepare_text_ids(x: torch.Tensor, t_coord: torch.Tensor | None = None):
- B, L, _ = x.shape
- out_ids = []
-
- for i in range(B):
- t = torch.arange(1) if t_coord is None else t_coord[i]
- h = torch.arange(1)
- w = torch.arange(1)
- l = torch.arange(L)
-
- coords = torch.cartesian_prod(t, h, w, l)
- out_ids.append(coords)
-
- return torch.stack(out_ids)
+ @property
+ def current_timestep(self):
+ return self._current_timestep
@staticmethod
- def _prepare_latent_ids(latents: torch.Tensor):
+ def _prepare_text_ids(x: torch.Tensor):
+ batch_size, sequence_length, _ = x.shape
+ text_ids = torch.cartesian_prod(
+ torch.arange(1), torch.arange(1), torch.arange(1), torch.arange(sequence_length)
+ )
+ return text_ids.unsqueeze(0).expand(batch_size, -1, -1)
+
+ @staticmethod
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._prepare_latent_ids
+ def _prepare_latent_ids(
+ latents: torch.Tensor, # (B, C, H, W)
+ ):
r"""
Generates 4D position coordinates (T, H, W, L) for latent tensors.
@@ -231,6 +224,7 @@ class SeFiPipeline(DiffusionPipeline):
return latent_ids
@staticmethod
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpatchify_latents
def _unpatchify_latents(latents):
batch_size, num_channels_latents, height, width = latents.shape
latents = latents.reshape(batch_size, num_channels_latents // (2 * 2), 2, 2, height, width)
@@ -239,6 +233,7 @@ class SeFiPipeline(DiffusionPipeline):
return latents
@staticmethod
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._pack_latents
def _pack_latents(latents):
"""
pack latents: (batch_size, num_channels, height, width) -> (batch_size, height * width, num_channels)
@@ -250,9 +245,8 @@ class SeFiPipeline(DiffusionPipeline):
return latents
@staticmethod
- def _unpack_latents_with_ids(
- x: torch.Tensor, x_ids: torch.Tensor, height: int | None = None, width: int | None = None
- ):
+ # Copied from diffusers.pipelines.flux2.pipeline_flux2.Flux2Pipeline._unpack_latents_with_ids
+ def _unpack_latents_with_ids(x: torch.Tensor, x_ids: torch.Tensor) -> list[torch.Tensor]:
"""
using position ids to scatter tokens into place
"""
@@ -306,34 +300,19 @@ class SeFiPipeline(DiffusionPipeline):
def _build_chat_text(self, prompt: str) -> str:
messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
- try:
- return self.tokenizer.apply_chat_template(
- messages,
- tokenize=False,
- add_generation_prompt=True,
- enable_thinking=False,
- )
- except TypeError:
- return self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ return self.tokenizer.apply_chat_template(
+ messages,
+ tokenize=False,
+ add_generation_prompt=True,
+ enable_thinking=False,
+ )
def _align_text_encoder_rotary_dtype(self, device: torch.device):
- text_encoder = self.text_encoder
- if text_encoder is None:
- return
-
- try:
- text_encoder_dtype = next(text_encoder.parameters()).dtype
- except StopIteration:
- return
-
- text_model = text_encoder.model if hasattr(text_encoder, "model") else text_encoder
- language_model = getattr(text_model, "language_model", None)
- rotary_emb = getattr(language_model, "rotary_emb", None)
- if rotary_emb is not None:
- # Qwen3-VL stores RoPE inverse frequencies as non-persistent buffers. `from_pretrained(torch_dtype=...)`
- # can leave them in fp32 even when text weights are bf16, while the reference SeFi wrapper casts the whole
- # text encoder module. Keep these buffers aligned before text encoding.
- rotary_emb.to(device=device, dtype=text_encoder_dtype)
+ text_encoder_dtype = next(self.text_encoder.parameters()).dtype
+ # Qwen3-VL stores RoPE inverse frequencies as non-persistent buffers. `from_pretrained(dtype=...)`
+ # can leave them in fp32 even when text weights are bf16, while the reference SeFi wrapper casts the whole
+ # text encoder module. Keep these buffers aligned before text encoding.
+ self.text_encoder.model.language_model.rotary_emb.to(device=device, dtype=text_encoder_dtype)
def _get_qwen3vl_prompt_embeds(
self,
@@ -390,10 +369,10 @@ class SeFiPipeline(DiffusionPipeline):
):
device = device or self._execution_device
dtype = dtype or (self.transformer.dtype if self.transformer is not None else self.text_encoder.dtype)
- max_sequence_length = max_sequence_length or self.config.max_sequence_length
- text_encoder_hidden_layers = text_encoder_hidden_layers or tuple(self.config.text_encoder_hidden_layers)
if prompt_embeds is None:
+ max_sequence_length = max_sequence_length or self.config.max_sequence_length
+ text_encoder_hidden_layers = text_encoder_hidden_layers or tuple(self.config.text_encoder_hidden_layers)
prompt_embeds = self._get_qwen3vl_prompt_embeds(
prompt=prompt,
device=device,
@@ -608,7 +587,7 @@ class SeFiPipeline(DiffusionPipeline):
dtype=torch.float32,
)
u_shifted_unit = _apply_timestep_shift_unit_interval(u_base_unit, self.config.timestep_shift_alpha)
- _, base_sigmas_schedule = self._timesteps_and_sigmas(u_shifted_unit, n_dim=1, dtype=torch.float32)
+ base_timesteps_schedule, _ = self._timesteps_and_sigmas(u_shifted_unit, n_dim=1, dtype=torch.float32)
u_sem_raw_schedule = u_shifted_unit * (1.0 + float(self.config.delta_t))
self._num_timesteps = num_inference_steps
@@ -631,9 +610,9 @@ class SeFiPipeline(DiffusionPipeline):
_, sigmas_sem_next = self._timesteps_and_sigmas(u_sem_next, latents.ndim, latents.dtype)
_, sigmas_tex_next = self._timesteps_and_sigmas(u_tex_next, latents.ndim, latents.dtype)
- self._current_timestep = base_sigmas_schedule[i]
+ self._current_timestep = base_timesteps_schedule[i]
packed_latents = self._pack_latents(latents)
- pred_cond = self.transformer(
+ noise_pred = self.transformer(
hidden_states=packed_latents,
timestep_sem=timesteps_sem_cur / 1000,
timestep_tex=timesteps_tex_cur / 1000,
@@ -643,8 +622,8 @@ class SeFiPipeline(DiffusionPipeline):
joint_attention_kwargs=self.attention_kwargs,
return_dict=False,
)[0]
- pred_cond = pred_cond[:, : packed_latents.size(1)]
- pred_cond = self._unpack_latents_with_ids(pred_cond, latent_ids)
+ noise_pred = noise_pred[:, : packed_latents.size(1)]
+ noise_pred = self._unpack_latents_with_ids(noise_pred, latent_ids)
if self.do_classifier_free_guidance:
pred_uncond = self.transformer(
@@ -659,9 +638,9 @@ class SeFiPipeline(DiffusionPipeline):
)[0]
pred_uncond = pred_uncond[:, : packed_latents.size(1)]
pred_uncond = self._unpack_latents_with_ids(pred_uncond, latent_ids)
- velocity = _combine_guided_velocity(pred_uncond, pred_cond, guidance_scale)
+ velocity = _combine_guided_velocity(pred_uncond, noise_pred, guidance_scale)
else:
- velocity = pred_cond
+ velocity = noise_pred
vel_sem = velocity[:, : self.semantic_channels]
vel_tex = velocity[:, self.semantic_channels :]
@@ -686,6 +665,8 @@ class SeFiPipeline(DiffusionPipeline):
progress_bar.update()
+ self._current_timestep = None
+
if output_type == "latent":
image = latents
else:
diff --git a/pipelines/sefi/transformer_sefi.py b/pipelines/sefi/transformer_sefi.py
index 96ec29309..75904c79c 100644
--- a/pipelines/sefi/transformer_sefi.py
+++ b/pipelines/sefi/transformer_sefi.py
@@ -19,10 +19,11 @@ import torch
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.utils import BaseOutput, apply_lora_scale
+from diffusers.utils import BaseOutput
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
+from diffusers.models.normalization import AdaLayerNormContinuous
+from diffusers.models.transformers.transformer_flux2 import Flux2Modulation, Flux2PosEmbed, Flux2SingleTransformerBlock, Flux2TransformerBlock
@dataclass
@@ -105,8 +106,6 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
RoPE theta.
eps (`float`, defaults to `1e-6`):
Normalization epsilon.
- text_input_dim (`int`, *optional*):
- Expected text embedding dimension. Defaults to `joint_attention_dim`.
"""
_supports_gradient_checkpointing = True
@@ -130,44 +129,60 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32),
rope_theta: int = 2000,
eps: float = 1e-6,
- text_input_dim: int | None = None,
):
super().__init__()
- text_input_dim = joint_attention_dim if text_input_dim is None else text_input_dim
- if int(text_input_dim) != int(joint_attention_dim):
- raise ValueError(
- f"`text_input_dim` must match `joint_attention_dim`, got {text_input_dim} and {joint_attention_dim}."
- )
-
self.out_channels = out_channels or in_channels
self.inner_dim = num_attention_heads * attention_head_dim
- self.backbone = Flux2Transformer2DModel(
- patch_size=patch_size,
- in_channels=in_channels,
- out_channels=out_channels,
- num_layers=num_layers,
- num_single_layers=num_single_layers,
- attention_head_dim=attention_head_dim,
- num_attention_heads=num_attention_heads,
- joint_attention_dim=joint_attention_dim,
- timestep_guidance_channels=timestep_guidance_channels,
- mlp_ratio=mlp_ratio,
- axes_dims_rope=axes_dims_rope,
- rope_theta=rope_theta,
- eps=eps,
- guidance_embeds=False,
- )
- # The reference SeFi transformer deletes Flux2's timestep/guidance embedder and stores only the dual embedder.
- self.backbone.time_guidance_embed = nn.Identity()
+
+ self.pos_embed = Flux2PosEmbed(theta=rope_theta, axes_dim=axes_dims_rope)
self.dual_time_embed = SeFiDualTimestepEmbeddings(
in_channels=timestep_guidance_channels,
embedding_dim=self.inner_dim,
bias=False,
)
+
+ self.double_stream_modulation_img = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+ self.double_stream_modulation_txt = Flux2Modulation(self.inner_dim, mod_param_sets=2, bias=False)
+ self.single_stream_modulation = Flux2Modulation(self.inner_dim, mod_param_sets=1, bias=False)
+
+ self.x_embedder = nn.Linear(in_channels, self.inner_dim, bias=False)
+ self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim, bias=False)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ Flux2TransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ mlp_ratio=mlp_ratio,
+ eps=eps,
+ bias=False,
+ )
+ for _ in range(num_layers)
+ ]
+ )
+ self.single_transformer_blocks = nn.ModuleList(
+ [
+ Flux2SingleTransformerBlock(
+ dim=self.inner_dim,
+ num_attention_heads=num_attention_heads,
+ attention_head_dim=attention_head_dim,
+ mlp_ratio=mlp_ratio,
+ eps=eps,
+ bias=False,
+ )
+ for _ in range(num_single_layers)
+ ]
+ )
+
+ self.norm_out = AdaLayerNormContinuous(
+ self.inner_dim, self.inner_dim, elementwise_affine=False, eps=eps, bias=False
+ )
+ self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=False)
+
self.gradient_checkpointing = False
- @apply_lora_scale("joint_attention_kwargs")
def forward(
self,
hidden_states: torch.Tensor,
@@ -211,26 +226,26 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
timestep_tex = timestep_tex.to(hidden_states.dtype) * 1000
temb = self.dual_time_embed(timestep_sem, timestep_tex)
- double_stream_mod_img = self.backbone.double_stream_modulation_img(temb)
- double_stream_mod_txt = self.backbone.double_stream_modulation_txt(temb)
- single_stream_mod = self.backbone.single_stream_modulation(temb)
+ double_stream_mod_img = self.double_stream_modulation_img(temb)
+ double_stream_mod_txt = self.double_stream_modulation_txt(temb)
+ single_stream_mod = self.single_stream_modulation(temb)
- hidden_states = self.backbone.x_embedder(hidden_states)
- encoder_hidden_states = self.backbone.context_embedder(encoder_hidden_states)
+ hidden_states = self.x_embedder(hidden_states)
+ encoder_hidden_states = self.context_embedder(encoder_hidden_states)
if img_ids.ndim == 3:
img_ids = img_ids[0]
if txt_ids.ndim == 3:
txt_ids = txt_ids[0]
- image_rotary_emb = self.backbone.pos_embed(img_ids)
- text_rotary_emb = self.backbone.pos_embed(txt_ids)
+ image_rotary_emb = self.pos_embed(img_ids)
+ text_rotary_emb = self.pos_embed(txt_ids)
concat_rotary_emb = (
torch.cat([text_rotary_emb[0], image_rotary_emb[0]], dim=0),
torch.cat([text_rotary_emb[1], image_rotary_emb[1]], dim=0),
)
- for block in self.backbone.transformer_blocks:
+ for block in self.transformer_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
block,
@@ -253,7 +268,7 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
- for block in self.backbone.single_transformer_blocks:
+ for block in self.single_transformer_blocks:
if torch.is_grad_enabled() and self.gradient_checkpointing:
hidden_states = self._gradient_checkpointing_func(
block,
@@ -273,8 +288,8 @@ class SeFiTransformer2DModel(ModelMixin, ConfigMixin):
)
hidden_states = hidden_states[:, num_txt_tokens:, ...]
- hidden_states = self.backbone.norm_out(hidden_states, temb)
- output = self.backbone.proj_out(hidden_states)
+ hidden_states = self.norm_out(hidden_states, temb)
+ output = self.proj_out(hidden_states)
if not return_dict:
return (output,)
diff --git a/pyproject.toml b/pyproject.toml
index 46f2a9a4a..f08cd1b64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -138,6 +138,7 @@ main.ignore-paths=[
"pipelines/hdm",
"pipelines/hidream",
"pipelines/lumina_dimmo",
+ "pipelines/mageflow",
"pipelines/meissonic",
"pipelines/omnigen2",
"pipelines/segmoe",
@@ -410,16 +411,16 @@ exclude = [
"pipelines/flex2",
"pipelines/hidream",
"pipelines/lumina_dimmo",
+ "pipelines/mageflow",
"pipelines/meissonic",
- "pipelines/meissonic/",
"pipelines/model_stablecascade.py",
- "pipelines/omnigen2/",
- "pipelines/sefi/",
- "pipelines/step1x/",
- "pipelines/ultraflux/",
- "pipelines/vibe/",
- "pipelines/xomni/",
- "pipelines/zetachroma/",
+ "pipelines/omnigen2",
+ "pipelines/sefi",
+ "pipelines/step1x",
+ "pipelines/ultraflux",
+ "pipelines/vibe",
+ "pipelines/xomni",
+ "pipelines/zetachroma",
"extensions-builtin/sd-extension-chainner/nodes",
]
diff --git a/requirements.txt b/requirements.txt
index 73b5e9ed7..27717958c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -17,21 +17,23 @@ voluptuous
fasteners
limits
orjson
-websockets
ftfy
+websockets
+httptools
# versioned
fastapi==0.124.4
+uvicorn==0.52.1
rich==15.0.0
safetensors==0.8.0
-peft==0.19.1
+peft==0.20.0
httpx==0.28.1
requests==2.34.2
-tqdm==4.68.3
+tqdm==4.70.0
accelerate==1.14.0
einops==0.8.2
-huggingface_hub==1.23.0
-hf_xet==1.5.1
+huggingface_hub==1.26.1
+hf_xet==1.6.0
numpy==2.1.2
pandas==2.3.1
protobuf==7.35.1
@@ -40,10 +42,11 @@ urllib3==1.26.19
Pillow==12.2.0
timm==1.0.27
pyparsing==3.3.2
-typing-extensions==4.15.0
+typing-extensions==4.16.0
sentencepiece==0.2.1
# lint
+ty
ruff
pylint
pre-commit
diff --git a/scripts/autocomplete.py b/scripts/autocomplete.py
index 664619afd..cea4dd92a 100644
--- a/scripts/autocomplete.py
+++ b/scripts/autocomplete.py
@@ -150,6 +150,7 @@ def on_update(selected):
class AutocompleteScript(scripts_manager.Script):
+ video_capable = scripts_manager.AlwaysVisible
def show(self, is_img2img):
return scripts_manager.AlwaysVisible
diff --git a/scripts/daam/utils.py b/scripts/daam/utils.py
index 58402c9ed..410de9b30 100644
--- a/scripts/daam/utils.py
+++ b/scripts/daam/utils.py
@@ -94,8 +94,8 @@ nlp = None
@lru_cache(maxsize=100000)
-def cached_nlp(prompt: str, type='en_core_web_md'):
- global nlp
+def cached_nlp(prompt: str, type='en_core_web_md'): # pylint: disable=redefined-builtin
+ global nlp # pylint: disable=global-statement
if nlp is None:
try:
diff --git a/scripts/nudenet_ext.py b/scripts/nudenet_ext.py
index d8c83f999..9d58c7eef 100644
--- a/scripts/nudenet_ext.py
+++ b/scripts/nudenet_ext.py
@@ -136,9 +136,10 @@ def process(
# defines script for dual-mode usage
+# see below for all available options and callbacks
+#
class ScriptNudeNet(scripts.Script):
- # see below for all available options and callbacks
- #
+ video_capable = scripts.AlwaysVisible
def title(self):
return 'NudeNet'
diff --git a/scripts/postprocessing_rembg.py b/scripts/postprocessing_rembg.py
index ca350b968..4068f958b 100644
--- a/scripts/postprocessing_rembg.py
+++ b/scripts/postprocessing_rembg.py
@@ -1,11 +1,11 @@
import os
import gradio as gr
from PIL import Image
-from modules import scripts_postprocessing
-
+from modules import errors, scripts_postprocessing
models = [
"none",
+ "lucida",
"ben2",
"silueta",
"u2net",
@@ -84,6 +84,15 @@ class ScriptPostprocessingRembg(scripts_postprocessing.ScriptPostprocessing):
image = ben2.remove(image, refine=refine)
except Exception as e:
log.error(f'RemoveBackground: model={model} {e}')
+ errors.display(e, 'Rembg model=ben2')
+ return pp
+ elif model == 'lucida':
+ try:
+ from modules.rembg import lucida
+ image = lucida.remove(image)
+ except Exception as e:
+ log.error(f'RemoveBackground: model={model} {e}')
+ errors.display(e, 'Rembg model=lucida')
return pp
else:
try:
@@ -101,6 +110,7 @@ class ScriptPostprocessingRembg(scripts_postprocessing.ScriptPostprocessing):
session=rembg.new_session(model))
except Exception as e:
log.error(f'RemoveBackground: model={model} {e}')
+ errors.display(e, f'Rembg model={model}')
return pp
if mask_only and image.mode == "RGBA":
diff --git a/scripts/postprocessing_seedvr.py b/scripts/postprocessing_seedvr.py
index 2796a3d46..0e47eb268 100644
--- a/scripts/postprocessing_seedvr.py
+++ b/scripts/postprocessing_seedvr.py
@@ -23,7 +23,7 @@ class ScriptSeedVR(scripts_postprocessing.ScriptPostprocessing):
seedvr_seed = gr.Number(step=1, value=-1, label="SeedVR seed", elem_id="extras_seedvr_seed")
with gr.Row():
seedvr_cfg_scale = gr.Slider(minimum=0.0, maximum=15.0, step=0.01, value=1.5, label="SeedVR guidance scale", elem_id="extras_seedvr_cfg_scale")
- seedvr_cfg_rescale = gr.Slider(minimum=0.0, maximum=15.0, step=0.01, value=0.0, label="SeedVR guidance rescale", elem_id="extras_seedvr_cfg_rescale")
+ seedvr_cfg_rescale = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.0, label="SeedVR guidance rescale", elem_id="extras_seedvr_cfg_rescale")
with gr.Accordion('SeedVR VAE', open = False, elem_id="postprocess_seedvr_vae_accordion"):
with gr.Row():
seedvr_vae_tile_encode = gr.Checkbox(label="VAE tiled encode", value=True, elem_id="extras_seedvr_vae_tile_encode")
@@ -32,8 +32,8 @@ class ScriptSeedVR(scripts_postprocessing.ScriptPostprocessing):
seedvr_tile_size = gr.Slider(minimum=64, maximum=4096, step=8, value=1024, label="SeedVR tile size", elem_id="extras_seedvr_tile_size")
seedvr_tile_overlap = gr.Slider(minimum=0, maximum=1.0, step=0.01, value=0.25, label="SeedVR tile overlap", elem_id="extras_seedvr_tile_overlap")
with gr.Row():
- seedvr_vae_memory = gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=1.0, label="SeedVR VAE memory", elem_id="extras_seedvr_vae_memory")
- with gr.Accordion('SeedVR video', open = False, elem_id="postprocess_seedvr_video_accordion"):
+ seedvr_vae_memory = gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=0.5, label="SeedVR VAE memory", elem_id="extras_seedvr_vae_memory")
+ with gr.Accordion('SeedVR video', open = False, elem_id="postprocess_seedvr_video_accordion", visible=False):
with gr.Row():
seedvr_batch_size = gr.Slider(minimum=1, maximum=64, step=1, value=1, label="SeedVR batch size", elem_id="extras_seedvr_batch_size")
seedvr_batch_overlap = gr.Slider(minimum=0, maximum=16, step=1, value=0, label="SeedVR batch overlap", elem_id="extras_seedvr_batch_overlap")
diff --git a/scripts/postprocessing_upscale.py b/scripts/postprocessing_upscale.py
index 92e6da29a..9726b23b3 100644
--- a/scripts/postprocessing_upscale.py
+++ b/scripts/postprocessing_upscale.py
@@ -68,7 +68,6 @@ class ScriptPostprocessingUpscale(scripts_postprocessing.ScriptPostprocessing):
return image
def process(self, pp: scripts_postprocessing.PostprocessedImage, upscale_mode=1, upscale_by=2.0, upscale_to_width=None, upscale_to_height=None, upscale_crop=False, upscaler_1_name=None, upscaler_2_name=None, upscaler_2_visibility=0.0): # pylint: disable=arguments-differ
-
if upscaler_1_name == "None":
upscaler_1_name = None
upscaler1 = next(iter([x for x in shared.sd_upscalers if x.name == upscaler_1_name]), None)
diff --git a/scripts/prompt_enhance/__init__.py b/scripts/prompt_enhance/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/scripts/prompt_enhance/helpers.py b/scripts/prompt_enhance/helpers.py
new file mode 100644
index 000000000..115372dac
--- /dev/null
+++ b/scripts/prompt_enhance/helpers.py
@@ -0,0 +1,88 @@
+import base64
+import io
+import gradio as gr
+from modules import ui_symbols
+from .options import Options
+
+
+def b64(image):
+ if image is None:
+ return ''
+ if isinstance(image, gr.Image): # should not happen
+ return None
+ with io.BytesIO() as stream:
+ image.convert('RGB').save(stream, 'JPEG')
+ values = stream.getvalue()
+ encoded = base64.b64encode(values).decode()
+ return encoded
+
+
+def is_cloud_model(model_name: str) -> bool:
+ if not model_name:
+ return False
+ return model_name in Options.cloud
+
+
+def is_vision_model(model_name: str) -> bool:
+ """Check if model supports vision/image input."""
+ if not model_name:
+ return False
+ return model_name in Options.img2img or model_name in Options.cloud
+
+
+def is_thinking_model(model_name: str) -> bool:
+ """Check if model supports thinking/reasoning mode."""
+ if not model_name:
+ return False
+ model_lower = model_name.lower()
+ # Match VQA's detection patterns for consistency
+ thinking_indicators = [
+ 'thinking', # Qwen3-VL-*-Thinking models
+ 'reasoning', # Ministral-3-*-Reasoning models
+ 'moondream3', # Moondream 3 supports thinking
+ 'moondream 3',
+ 'moondream2', # Moondream 2 supports reasoning mode
+ 'moondream 2',
+ 'mimo', # XiaomiMiMo models
+ 'qwen3.5', # Qwen3.5 native thinking (repo names)
+ 'qwen 3.5', # Qwen3.5 native thinking (display names)
+ ]
+ return any(indicator in model_lower for indicator in thinking_indicators)
+
+
+def get_model_display_name(model_repo: str) -> str:
+ """Generate display name with vision/reasoning symbols."""
+ symbols = []
+ if model_repo in Options.img2img:
+ symbols.append(ui_symbols.vision)
+ if model_repo in Options.cloud:
+ symbols.append(ui_symbols.cloud)
+ if is_thinking_model(model_repo):
+ symbols.append(ui_symbols.reasoning)
+ return f"{model_repo} {' '.join(symbols)}" if symbols else model_repo
+
+
+def get_model_repo_from_display(display_name: str) -> str:
+ """Strip symbols from display name to get repo."""
+ if not display_name:
+ return display_name
+ result = display_name
+ for symbol in [ui_symbols.vision, ui_symbols.reasoning, ui_symbols.cloud]:
+ result = result.replace(symbol, '')
+ return result.strip()
+
+
+def keep_think_block_open(text_prompt: str) -> str:
+ """Remove closing so model can continue reasoning with prefill."""
+ think_open = ""
+ think_close = ""
+ last_open = text_prompt.rfind(think_open)
+ if last_open == -1:
+ return text_prompt
+ close_index = text_prompt.find(think_close, last_open)
+ if close_index == -1:
+ return text_prompt
+ end_close = close_index + len(think_close)
+ while end_close < len(text_prompt) and text_prompt[end_close] in ' \t\r\n':
+ end_close += 1
+ return text_prompt[:close_index] + text_prompt[end_close:]
diff --git a/scripts/prompt_enhance/options.py b/scripts/prompt_enhance/options.py
new file mode 100644
index 000000000..1913de1bb
--- /dev/null
+++ b/scripts/prompt_enhance/options.py
@@ -0,0 +1,223 @@
+from dataclasses import dataclass
+import textwrap
+import transformers
+
+
+@dataclass
+class Options:
+ img2img = [
+ # Gemma
+ 'google/gemma-3-4b-it',
+ 'google/gemma-3n-E2B-it',
+ 'google/gemma-3n-E4B-it',
+ 'google/gemma-4-E2B-it',
+ 'google/gemma-4-E4B-it',
+ 'google/gemma-4-12B-it-qat-w4a16-ct',
+ # Qwen3.5
+ 'Qwen/Qwen3.5-2B',
+ 'Qwen/Qwen3.5-4B',
+ 'Qwen/Qwen3.5-9B',
+ # Qwen3-VL
+ 'Qwen/Qwen3-VL-2B-Instruct',
+ 'Qwen/Qwen3-VL-2B-Thinking',
+ 'Qwen/Qwen3-VL-4B-Instruct',
+ 'Qwen/Qwen3-VL-4B-Thinking',
+ 'Qwen/Qwen3-VL-8B-Instruct',
+ 'Qwen/Qwen3-VL-8B-Thinking',
+ # Qwen2.5-VL
+ 'Qwen/Qwen2.5-VL-3B-Instruct',
+ # Mistral
+ 'mistralai/Ministral-3-3B-Instruct-2512-BF16',
+ 'mistralai/Ministral-3-8B-Instruct-2512-BF16',
+ 'mistralai/Ministral-3-3B-Reasoning-2512',
+ 'mistralai/Ministral-3-8B-Reasoning-2512',
+ # Finetunes
+ 'trohrbaugh/gemma-4-E4B-it-heretic-ara',
+ 'trohrbaugh/Qwen3.5-9B-heretic-v2',
+ ]
+ cloud = [
+ 'google/gemini-3.5-flash',
+ 'google/gemini-3.1-pro-preview',
+ 'google/gemini-3.1-flash-lite',
+ 'google/gemini-3.1-flash-lite-preview',
+ 'google/gemini-2.5-flash',
+ 'google/gemini-2.5-flash-lite',
+ 'google/gemini-2.5-pro',
+ ]
+ models = {
+ # Gemma
+ 'google/gemma-3-1b-it': {},
+ 'google/gemma-3-4b-it': {},
+ 'google/gemma-3n-E2B-it': {},
+ 'google/gemma-3n-E4B-it': {},
+ 'google/gemma-4-E2B-it': {},
+ 'google/gemma-4-E4B-it': {},
+ 'google/gemma-4-12B-it-qat-w4a16-ct': {}, # compressed-tensor model
+ # Qwen3.5
+ 'Qwen/Qwen3.5-0.8B': {},
+ 'Qwen/Qwen3.5-2B': {},
+ 'Qwen/Qwen3.5-4B': {},
+ 'Qwen/Qwen3.5-9B': {},
+ # Qwen3
+ 'Qwen/Qwen3-0.6B': {},
+ 'Qwen/Qwen3-1.7B': {},
+ 'Qwen/Qwen3-4B': {},
+ 'Qwen/Qwen3-4B-Instruct-2507': {},
+ # Qwen3-VL
+ 'Qwen/Qwen3-VL-2B-Instruct': {},
+ 'Qwen/Qwen3-VL-2B-Thinking': {},
+ 'Qwen/Qwen3-VL-4B-Instruct': {},
+ 'Qwen/Qwen3-VL-4B-Thinking': {},
+ 'Qwen/Qwen3-VL-8B-Instruct': {},
+ 'Qwen/Qwen3-VL-8B-Thinking': {},
+ # Qwen2.5
+ 'Qwen/Qwen2.5-0.5B-Instruct': {},
+ 'Qwen/Qwen2.5-1.5B-Instruct': {},
+ 'Qwen/Qwen2.5-3B-Instruct': {},
+ # Qwen2.5-VL
+ 'Qwen/Qwen2.5-VL-3B-Instruct': {},
+ # Llama
+ 'meta-llama/Llama-3.2-1B-Instruct': {},
+ 'meta-llama/Llama-3.2-3B-Instruct': {},
+ 'meta-llama/Llama-3.2-8B-Instruct': {},
+ 'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {},
+ 'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {},
+ # Gemini
+ 'google/gemini-3.5-flash': {},
+ 'google/gemini-3.1-pro-preview': {},
+ 'google/gemini-3.1-flash-lite': {},
+ 'google/gemini-3.1-flash-lite-preview': {},
+ 'google/gemini-2.5-flash': {},
+ 'google/gemini-2.5-flash-lite': {},
+ 'google/gemini-2.5-pro': {},
+ # SmolLM
+ 'HuggingFaceTB/SmolLM2-135M-Instruct': {},
+ 'HuggingFaceTB/SmolLM2-360M-Instruct': {},
+ 'HuggingFaceTB/SmolLM2-1.7B-Instruct': {},
+ 'HuggingFaceTB/SmolLM3-3B': {},
+ # Phi
+ 'microsoft/Phi-4-mini-instruct': {},
+ # Mistral
+ 'mistralai/Ministral-3-3B-Instruct-2512-BF16': {},
+ 'mistralai/Ministral-3-8B-Instruct-2512-BF16': {},
+ 'mistralai/Ministral-3-3B-Reasoning-2512': {},
+ 'mistralai/Ministral-3-8B-Reasoning-2512': {},
+ # Finetunes
+ 'p-e-w/gemma-4-E2B-it-heretic-ara': {},
+ 'trohrbaugh/gemma-4-E4B-it-heretic-ara': {},
+ 'trohrbaugh/Qwen3.5-9B-heretic-v2': {},
+ # GGUF
+ 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF': { # kept primarily as an example how to add gguf model
+ 'repo': 'meta-llama/Llama-3.2-1B-Instruct', # original repo so we can load missing components
+ 'type': 'llama', # required so gguf loader knows what to do
+ 'gguf': 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF', # gguf repo
+ 'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo
+ },
+ }
+ models_cls = {
+ 'qwen3_5': 'Qwen3_5ForConditionalGeneration',
+ 'qwen3_5_moe': 'Qwen3_5MoeForConditionalGeneration',
+ 'qwen3_vl': 'Qwen3VLForConditionalGeneration',
+ 'qwen2_5_vl': 'Qwen2_5_VLForConditionalGeneration',
+ 'qwen2_vl': 'Qwen2VLForConditionalGeneration',
+ 'mistral3': 'Mistral3ForConditionalGeneration',
+ 'gemma4': 'Gemma4ForConditionalGeneration',
+ }
+
+ # default = list(models)[1] # gemma-3-4b-it
+ default = 'google/gemma-3-4b-it'
+ supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING)
+ t2i_prompt: str = textwrap.dedent('''\
+ You are an expert AI image prompt engineer.
+ You will receive a user prompt for image generation.
+ Your sole job is to rewrite user inputs into highly detailed, visually rich prompts for image generation models.
+ Improve the prompt by adding relevant visual specificity for composition, lighting, color, texture, and atmosphere.
+ Keep the result faithful to the original prompt and the intended image.
+ Do not add unrelated concepts, non-visual commentary, or fluff.
+ ''')
+ i2i_prompt: str = textwrap.dedent('''\
+ You are an expert AI image prompt engineer.
+ You will receive an image and a user prompt for editing or refinement.
+ Your sole job is to rewrite user inputs into highly detailed, visually rich prompts for image generation models while taking the provided image into account.
+ Improve the prompt with concrete visual detail that remains faithful to the image and edit intent.
+ Keep the result grounded in image-generation language.
+ Do not invent unrelated objects, actions, or concepts.
+ ''')
+ i2i_noprompt: str = textwrap.dedent('''\
+ You are an expert AI image prompt engineer.
+ You will receive an image and must create an editing prompt from the visible content.
+ Your sole job is to create highly detailed, visually rich prompts for image generation models based on the provided image.
+ Base the prompt on what is present in the image, its style, and its atmosphere.
+ Use concrete, image-generation language and avoid unrelated fluff or commentary.
+ ''')
+ t2v_prompt: str = textwrap.dedent('''\
+ You are an expert AI video prompt engineer.
+ You will receive a user prompt for video generation.
+ Your sole job is to rewrite user inputs into highly detailed, visually rich prompts for video generation models.
+ Improve the prompt by adding relevant visual specificity for composition, lighting, color, texture, and atmosphere.
+ Specify subject and camera movement, scene transitions, and temporal dynamics.
+ Keep the result faithful to the original prompt and the intended video.
+ Do not add unrelated concepts, non-visual commentary, or fluff.
+ ''')
+ i2v_prompt: str = textwrap.dedent('''\
+ You are an expert AI video prompt engineer.
+ You will receive an image and a user prompt for editing or refinement.
+ Your sole job is to rewrite user inputs into highly detailed, visually rich prompts for video generation models while taking the provided image into account.
+ Improve the prompt with concrete visual detail that remains faithful to the image and edit intent.
+ Specify subject and camera movement, scene transitions, and temporal dynamics.
+ Keep the result grounded in video-generation language.
+ Do not invent unrelated objects, actions, or concepts.
+ ''')
+ i2v_noprompt: str = textwrap.dedent('''\
+ You are an expert AI video prompt engineer.
+ You will receive an image and must create a video prompt from the visible content.
+ Your sole job is to create highly detailed, visually rich prompts for video generation models based on the provided image.
+ Base the prompt on what is present in the image, its style, and its atmosphere.
+ Specify subject and camera movement, scene transitions, and temporal dynamics.
+ Use concrete, video-generation language and avoid unrelated fluff or commentary.
+ ''')
+ nsfw_ok: str = textwrap.dedent('''\
+ NSFW, nudity and explicit content and language are allowed.
+ ''')
+ nsfw_no: str = textwrap.dedent('''\
+ NSFW, nudity or explicit content and language are not allowed; if present, remove those elements from the enhanced prompt.
+ ''')
+ details_prompt: str = textwrap.dedent('''\
+ Enhance the prompt with relevant image details.
+ Identify the main subject(s), physical attributes, clothing if non-nude, and distinctive visual features.
+ Describe the scene, environment, key objects, and overall mood or atmosphere.
+ Be precise and efficient; avoid redundancy, abstract commentary, unrelated fluff, or instructions.
+ Do not invent any objects, settings, or themes not implied by the input.
+ Do not add era, background props, or atmosphere unless explicitly present in the prompt.
+ ''')
+ details_format: str = textwrap.dedent('''\
+ Output exactly one enhanced prompt string.
+ Do not add greetings, comments, explanations, follow-up questions, labels, formatting, or numbering.
+ Do not include any extra prose or analysis.
+ Start immediately with the prompt content.
+ No stray tokens!
+ ''')
+
+ censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed']
+
+ max_delim_index: int = 60
+ min_tokens: int = 0
+ max_tokens: int = 256
+ do_sample: bool = True
+ temperature: float = 0.6
+ repetition_penalty: float = 1.2
+ top_k: int = 0
+ top_p: float = 0.0
+ thinking_mode: bool = False
+
+ @staticmethod
+ def get_model_choices():
+ """Return list of display names for dropdown."""
+ from .helpers import get_model_display_name
+ return [get_model_display_name(repo) for repo in Options.models.keys()]
+
+ @staticmethod
+ def get_default_display():
+ """Return display name for default model."""
+ from .helpers import get_model_display_name
+ return get_model_display_name(Options.default)
diff --git a/scripts/prompt_enhance/template.py b/scripts/prompt_enhance/template.py
new file mode 100644
index 000000000..365508bf3
--- /dev/null
+++ b/scripts/prompt_enhance/template.py
@@ -0,0 +1,100 @@
+import os
+from PIL import Image
+from modules.logger import log
+from .options import Options
+from .helpers import b64, is_cloud_model
+
+
+debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None
+debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
+
+
+def get_text_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, _image) -> list[dict]:
+ if not has_system:
+ system = options.t2v_prompt if is_video else options.t2i_prompt
+ system += options.nsfw_ok if nsfw else options.nsfw_no
+ system += options.details_prompt
+ system += options.details_format
+ debug_log(f'Prompt enhance: system="{system}"')
+ if not has_prompt:
+ prompt = 'be creative!'
+ if not has_processor:
+ chat_template = [
+ { "role": "system", "content": system },
+ { "role": "user", "content": prompt },
+ ]
+ else:
+ chat_template = [
+ { "role": "system", "content": [
+ {"type": "text", "text": system }
+ ] },
+ { "role": "user", "content": [
+ {"type": "text", "text": prompt},
+ ] },
+ ]
+ return chat_template
+
+
+def get_image_template(system, prompt, options, nsfw, has_system, has_prompt, _has_processor, is_video, image) -> list[dict]:
+ if not has_system:
+ if is_video:
+ system = options.i2v_prompt if has_prompt else options.i2v_noprompt
+ else:
+ system = options.i2i_prompt if has_prompt else options.i2i_noprompt
+ system += options.nsfw_ok if nsfw else options.nsfw_no
+ system += options.details_prompt
+ system += options.details_format
+ debug_log(f'Prompt enhance: system="{system}"')
+ if has_prompt:
+ chat_template = [
+ { "role": "system", "content": [
+ {"type": "text", "text": system }
+ ] },
+ { "role": "user", "content": [
+ {"type": "text", "text": prompt},
+ {"type": "image", "image": b64(image)}
+ ] },
+ ]
+ else:
+ chat_template = [
+ { "role": "system", "content": [
+ {"type": "text", "text": system }
+ ] },
+ { "role": "user", "content": [
+ {"type": "image", "image": b64(image)}
+ ] },
+ ]
+ return chat_template
+
+
+def set_template(
+ system: str | None,
+ prompt: str | None,
+ image: Image.Image | None,
+ options: Options,
+ model: str,
+ nsfw: bool = True,
+ has_processor: bool = False,
+ module: str | None = None,
+) -> list[dict] | str:
+ chat_template = []
+ has_system = system is not None and len(system) > 4
+ has_prompt = prompt is not None and len(prompt) > 4
+ has_image = image is not None and isinstance(image, Image.Image)
+ is_video = module == 'video'
+
+ debug_log(f'Prompt enhance template: module={module} system={has_system} prompt={has_prompt} image={has_image} video={is_video} model="{model}" nsfw={nsfw} processor={has_processor}')
+
+ if has_image:
+ if is_cloud_model(model):
+ pass
+ elif options.processor is None:
+ log.error('Prompt enhance: image not supported by model')
+ return prompt # Return original text part if image cannot be processed
+
+ if has_image:
+ chat_template = get_image_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, image)
+ else:
+ chat_template = get_text_template(system, prompt, options, nsfw, has_system, has_prompt, has_processor, is_video, image)
+
+ return chat_template
diff --git a/scripts/prompt_enhance.py b/scripts/prompt_enhance_ext.py
similarity index 73%
rename from scripts/prompt_enhance.py
rename to scripts/prompt_enhance_ext.py
index c00357b71..14d801359 100644
--- a/scripts/prompt_enhance.py
+++ b/scripts/prompt_enhance_ext.py
@@ -1,299 +1,26 @@
-from dataclasses import dataclass
-import io
import os
import re
import time
import random
-import base64
-import textwrap
import torch
import transformers
import gradio as gr
from PIL import Image
-from modules import scripts_manager, shared, devices, errors, processing, sd_models, sd_modules, timer, ui_symbols
+from modules import scripts_manager, shared, devices, errors, processing, sd_models, sd_modules, timer
from modules import ui_control_helpers
from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
from modules.logger import log
from modules.caption.logits import LogitsParser
from modules.caption import helpers
+from scripts.prompt_enhance.options import Options
+from scripts.prompt_enhance.helpers import is_cloud_model, is_vision_model, is_thinking_model, get_model_repo_from_display
+from scripts.prompt_enhance.template import set_template
debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None
debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None
-def b64(image):
- if image is None:
- return ''
- if isinstance(image, gr.Image): # should not happen
- return None
- with io.BytesIO() as stream:
- image.convert('RGB').save(stream, 'JPEG')
- values = stream.getvalue()
- encoded = base64.b64encode(values).decode()
- return encoded
-
-
-def is_cloud_model(model_name: str) -> bool:
- if not model_name:
- return False
- return model_name in Options.cloud
-
-
-def is_vision_model(model_name: str) -> bool:
- """Check if model supports vision/image input."""
- if not model_name:
- return False
- return model_name in Options.img2img or model_name in Options.cloud
-
-
-def is_thinking_model(model_name: str) -> bool:
- """Check if model supports thinking/reasoning mode."""
- if not model_name:
- return False
- model_lower = model_name.lower()
- # Match VQA's detection patterns for consistency
- thinking_indicators = [
- 'thinking', # Qwen3-VL-*-Thinking models
- 'reasoning', # Ministral-3-*-Reasoning models
- 'moondream3', # Moondream 3 supports thinking
- 'moondream 3',
- 'moondream2', # Moondream 2 supports reasoning mode
- 'moondream 2',
- 'mimo', # XiaomiMiMo models
- 'qwen3.5', # Qwen3.5 native thinking (repo names)
- 'qwen 3.5', # Qwen3.5 native thinking (display names)
- ]
- return any(indicator in model_lower for indicator in thinking_indicators)
-
-
-def get_model_display_name(model_repo: str) -> str:
- """Generate display name with vision/reasoning symbols."""
- symbols = []
- if model_repo in Options.img2img:
- symbols.append(ui_symbols.vision)
- if model_repo in Options.cloud:
- symbols.append(ui_symbols.cloud)
- if is_thinking_model(model_repo):
- symbols.append(ui_symbols.reasoning)
- return f"{model_repo} {' '.join(symbols)}" if symbols else model_repo
-
-
-def get_model_repo_from_display(display_name: str) -> str:
- """Strip symbols from display name to get repo."""
- if not display_name:
- return display_name
- result = display_name
- for symbol in [ui_symbols.vision, ui_symbols.reasoning, ui_symbols.cloud]:
- result = result.replace(symbol, '')
- return result.strip()
-
-
-def keep_think_block_open(text_prompt: str) -> str:
- """Remove closing so model can continue reasoning with prefill."""
- think_open = ""
- think_close = ""
- last_open = text_prompt.rfind(think_open)
- if last_open == -1:
- return text_prompt
- close_index = text_prompt.find(think_close, last_open)
- if close_index == -1:
- return text_prompt
- end_close = close_index + len(think_close)
- while end_close < len(text_prompt) and text_prompt[end_close] in ' \t\r\n':
- end_close += 1
- return text_prompt[:close_index] + text_prompt[end_close:]
-
-
-@dataclass
-class Options:
- img2img = [
- # Gemma
- 'google/gemma-3-4b-it',
- 'google/gemma-3n-E2B-it',
- 'google/gemma-3n-E4B-it',
- 'google/gemma-4-E2B-it',
- 'google/gemma-4-E4B-it',
- 'google/gemma-4-12B-it-qat-w4a16-ct',
- # Qwen3.5
- 'Qwen/Qwen3.5-2B',
- 'Qwen/Qwen3.5-4B',
- 'Qwen/Qwen3.5-9B',
- # Qwen3-VL
- 'Qwen/Qwen3-VL-2B-Instruct',
- 'Qwen/Qwen3-VL-2B-Thinking',
- 'Qwen/Qwen3-VL-4B-Instruct',
- 'Qwen/Qwen3-VL-4B-Thinking',
- 'Qwen/Qwen3-VL-8B-Instruct',
- 'Qwen/Qwen3-VL-8B-Thinking',
- # Qwen2.5-VL
- 'Qwen/Qwen2.5-VL-3B-Instruct',
- # Mistral
- 'mistralai/Ministral-3-3B-Instruct-2512-BF16',
- 'mistralai/Ministral-3-8B-Instruct-2512-BF16',
- 'mistralai/Ministral-3-3B-Reasoning-2512',
- 'mistralai/Ministral-3-8B-Reasoning-2512',
- # Finetunes
- 'trohrbaugh/gemma-4-E4B-it-heretic-ara',
- 'trohrbaugh/Qwen3.5-9B-heretic-v2',
- ]
- cloud = [
- 'google/gemini-3.5-flash',
- 'google/gemini-3.1-pro-preview',
- 'google/gemini-3.1-flash-lite',
- 'google/gemini-3.1-flash-lite-preview',
- 'google/gemini-2.5-flash',
- 'google/gemini-2.5-flash-lite',
- 'google/gemini-2.5-pro',
- ]
- models = {
- # Gemma
- 'google/gemma-3-1b-it': {},
- 'google/gemma-3-4b-it': {},
- 'google/gemma-3n-E2B-it': {},
- 'google/gemma-3n-E4B-it': {},
- 'google/gemma-4-E2B-it': {},
- 'google/gemma-4-E4B-it': {},
- 'google/gemma-4-12B-it-qat-w4a16-ct': {}, # compressed-tensor model
- # Qwen3.5
- 'Qwen/Qwen3.5-0.8B': {},
- 'Qwen/Qwen3.5-2B': {},
- 'Qwen/Qwen3.5-4B': {},
- 'Qwen/Qwen3.5-9B': {},
- # Qwen3
- 'Qwen/Qwen3-0.6B': {},
- 'Qwen/Qwen3-1.7B': {},
- 'Qwen/Qwen3-4B': {},
- 'Qwen/Qwen3-4B-Instruct-2507': {},
- # Qwen3-VL
- 'Qwen/Qwen3-VL-2B-Instruct': {},
- 'Qwen/Qwen3-VL-2B-Thinking': {},
- 'Qwen/Qwen3-VL-4B-Instruct': {},
- 'Qwen/Qwen3-VL-4B-Thinking': {},
- 'Qwen/Qwen3-VL-8B-Instruct': {},
- 'Qwen/Qwen3-VL-8B-Thinking': {},
- # Qwen2.5
- 'Qwen/Qwen2.5-0.5B-Instruct': {},
- 'Qwen/Qwen2.5-1.5B-Instruct': {},
- 'Qwen/Qwen2.5-3B-Instruct': {},
- # Qwen2.5-VL
- 'Qwen/Qwen2.5-VL-3B-Instruct': {},
- # Llama
- 'meta-llama/Llama-3.2-1B-Instruct': {},
- 'meta-llama/Llama-3.2-3B-Instruct': {},
- 'meta-llama/Llama-3.2-8B-Instruct': {},
- 'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {},
- 'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {},
- # Gemini
- 'google/gemini-3.5-flash': {},
- 'google/gemini-3.1-pro-preview': {},
- 'google/gemini-3.1-flash-lite': {},
- 'google/gemini-3.1-flash-lite-preview': {},
- 'google/gemini-2.5-flash': {},
- 'google/gemini-2.5-flash-lite': {},
- 'google/gemini-2.5-pro': {},
- # SmolLM
- 'HuggingFaceTB/SmolLM2-135M-Instruct': {},
- 'HuggingFaceTB/SmolLM2-360M-Instruct': {},
- 'HuggingFaceTB/SmolLM2-1.7B-Instruct': {},
- 'HuggingFaceTB/SmolLM3-3B': {},
- # Phi
- 'microsoft/Phi-4-mini-instruct': {},
- # Mistral
- 'mistralai/Ministral-3-3B-Instruct-2512-BF16': {},
- 'mistralai/Ministral-3-8B-Instruct-2512-BF16': {},
- 'mistralai/Ministral-3-3B-Reasoning-2512': {},
- 'mistralai/Ministral-3-8B-Reasoning-2512': {},
- # Finetunes
- 'p-e-w/gemma-4-E2B-it-heretic-ara': {},
- 'trohrbaugh/gemma-4-E4B-it-heretic-ara': {},
- 'trohrbaugh/Qwen3.5-9B-heretic-v2': {},
- # GGUF
- 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF': { # kept primarily as an example how to add gguf model
- 'repo': 'meta-llama/Llama-3.2-1B-Instruct', # original repo so we can load missing components
- 'type': 'llama', # required so gguf loader knows what to do
- 'gguf': 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF', # gguf repo
- 'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo
- },
- }
- models_cls = {
- 'qwen3_5': 'Qwen3_5ForConditionalGeneration',
- 'qwen3_5_moe': 'Qwen3_5MoeForConditionalGeneration',
- 'qwen3_vl': 'Qwen3VLForConditionalGeneration',
- 'qwen2_5_vl': 'Qwen2_5_VLForConditionalGeneration',
- 'qwen2_vl': 'Qwen2VLForConditionalGeneration',
- 'mistral3': 'Mistral3ForConditionalGeneration',
- 'gemma4': 'Gemma4ForConditionalGeneration',
- }
-
- # default = list(models)[1] # gemma-3-4b-it
- default = 'google/gemma-3-4b-it'
- supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING)
- t2i_prompt: str = textwrap.dedent('''\
- You are an image prompt expert.
- You will receive a user prompt for image generation.
- Improve the prompt by adding relevant visual specificity for composition, lighting, color, texture, and atmosphere.
- Keep the result faithful to the original prompt and the intended image.
- Do not add unrelated concepts, non-visual commentary, or fluff.
- ''')
- i2i_prompt: str = textwrap.dedent('''\
- You are an image prompt expert.
- You will receive an image and a user prompt for editing or refinement.
- Improve the prompt with concrete visual detail that remains faithful to the image and edit intent.
- Keep the result grounded in image-generation language.
- Do not invent unrelated objects, actions, or concepts.
- ''')
- i2i_noprompt: str = textwrap.dedent('''\
- You are an image prompt expert.
- You will receive an image and must create an editing prompt from the visible content.
- Base the prompt on what is present in the image, its style, and its atmosphere.
- Use concrete, image-generation language and avoid unrelated fluff or commentary.
- ''')
- nsfw_ok: str = textwrap.dedent('''\
- NSFW, nudity and explicit content and language are allowed.
- ''')
- nsfw_no: str = textwrap.dedent('''\
- NSFW, nudity or explicit content and language are not allowed; if present, remove those elements from the enhanced prompt.
- ''')
- details_prompt: str = textwrap.dedent('''\
- Enhance the prompt with relevant image details.
- Identify the main subject(s), physical attributes, clothing if non-nude, and distinctive visual features.
- Describe the scene, environment, key objects, and overall mood or atmosphere.
- Be precise and efficient; avoid redundancy, abstract commentary, unrelated fluff, or instructions.
- Do not invent any objects, settings, or themes not implied by the input.
- Do not add era, background props, or atmosphere unless explicitly present in the prompt.
- ''')
- details_format: str = textwrap.dedent('''\
- Output exactly one enhanced prompt string.
- Do not add greetings, comments, explanations, follow-up questions, labels, formatting, or numbering.
- Do not include any extra prose or analysis.
- Start immediately with the prompt content.
- No stray tokens!
- ''')
-
- censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed']
-
- max_delim_index: int = 60
- min_tokens: int = 0
- max_tokens: int = 256
- do_sample: bool = True
- temperature: float = 0.6
- repetition_penalty: float = 1.2
- top_k: int = 0
- top_p: float = 0.0
- thinking_mode: bool = False
-
- @staticmethod
- def get_model_choices():
- """Return list of display names for dropdown."""
- return [get_model_display_name(repo) for repo in Options.models.keys()]
-
- @staticmethod
- def get_default_display():
- """Return display name for default model."""
- return get_model_display_name(Options.default)
-
-
class PromptEnhanceScript(scripts_manager.Script):
prompt: gr.Textbox = None
image: gr.Image = None
@@ -304,6 +31,7 @@ class PromptEnhanceScript(scripts_manager.Script):
busy: bool = False
server = None
options = Options()
+ video_capable = scripts_manager.AlwaysVisible
def title(self):
return 'Prompt enhance'
@@ -352,7 +80,7 @@ class PromptEnhanceScript(scripts_manager.Script):
gguf_args['model_type'] = model_type
gguf_args['gguf_file'] = model_file
- quant_args = model_quant.create_config(module='LLM') if not gguf_args else {}
+ quant_args = model_quant.create_config(module='LLM', modules_to_not_convert=['conv1d', 'linear_attn.conv1d']) if not gguf_args else {}
try:
t0 = time.time()
@@ -399,6 +127,7 @@ class PromptEnhanceScript(scripts_manager.Script):
)
finally:
sd_models.set_huggingface_options(quiet=True)
+
self.llm.eval()
register_aux('prompt_enhance', self.llm)
tokenizer_args = { 'pretrained_model_name_or_path': model_repo }
@@ -414,12 +143,11 @@ class PromptEnhanceScript(scripts_manager.Script):
debug_log(f'Prompt enhance: {m}')
self.model = name
t1 = time.time()
- log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} name="{name}" repo="{model_repo}" fn="{model_file}" processor="{self.processor.__class__.__name__ if self.processor else None}" tokenizer="{self.tokenizer.__class__.__name__ if self.tokenizer else None}" time={t1-t0:.2f} loaded')
+ log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} name="{name}" repo="{model_repo}" fn="{model_file}" processor="{self.processor.__class__.__name__ if self.processor else None}" tokenizer="{self.tokenizer.__class__.__name__ if self.tokenizer else None}" module={self.parent} time={t1-t0:.2f} loaded')
self.compile()
except Exception as e:
log.error(f'Prompt enhance: load {e}')
- if debug_enabled:
- errors.display(e, 'Prompt enhance')
+ errors.display(e, 'Prompt enhance')
devices.torch_gc()
self.set_openai(enable=use_openai)
@@ -661,63 +389,16 @@ class PromptEnhanceScript(scripts_manager.Script):
current_image = current_image.convert('RGB')
debug_log('Prompt enhance: Converted image to RGB mode')
- has_system = system is not None and len(system) > 4
-
- if current_image is not None and isinstance(current_image, Image.Image):
- if is_cloud_model(model):
- pass
- elif self.processor is None:
- log.error('Prompt enhance: image not supported by model')
- return prompt_text # Return original text part if image cannot be processed
- if prompt_text is not None and len(prompt_text) > 0:
- if not has_system:
- system = self.options.i2i_prompt
- system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
- system += self.options.details_prompt
- system += self.options.details_format
- chat_template = [
- { "role": "system", "content": [
- {"type": "text", "text": system }
- ] },
- { "role": "user", "content": [
- {"type": "text", "text": prompt_text},
- {"type": "image", "image": b64(current_image)}
- ] },
- ]
- else:
- if not has_system:
- system = self.options.i2i_noprompt
- system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
- system += self.options.details_prompt
- system += self.options.details_format
- chat_template = [
- { "role": "system", "content": [
- {"type": "text", "text": system }
- ] },
- { "role": "user", "content": [
- {"type": "image", "image": b64(current_image)}
- ] },
- ]
- else:
- if not has_system:
- system = self.options.t2i_prompt
- system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
- system += self.options.details_prompt
- system += self.options.details_format
- if self.processor is None:
- chat_template = [
- { "role": "system", "content": system },
- { "role": "user", "content": prompt_text },
- ]
- else:
- chat_template = [
- { "role": "system", "content": [
- {"type": "text", "text": system }
- ] },
- { "role": "user", "content": [
- {"type": "text", "text": prompt_text},
- ] },
- ]
+ chat_template = set_template(
+ system=system,
+ prompt=prompt_text,
+ image=current_image,
+ options=self.options,
+ model=model,
+ nsfw=nsfw,
+ has_processor=self.processor is not None,
+ module=self.parent,
+ )
# Prepare prefill (VQA approach: string concatenation, not assistant message)
prefill_text = (prefill or '').strip()
@@ -809,7 +490,7 @@ class PromptEnhanceScript(scripts_manager.Script):
return prompt_text # Return original text part on error
try:
- with devices.inference_context():
+ with devices.llm_context():
move_aux_to_gpu('prompt_enhance')
gen_kwargs = {
'do_sample': sample,
@@ -835,7 +516,7 @@ class PromptEnhanceScript(scripts_manager.Script):
log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} model="{model}" tokens={input_len} args={gen_kwargs} custom={custom}')
defaults = {k: v for k, v in helpers.get_default_args(self.llm).items() if k not in gen_kwargs}
- log.debug(f'Prompt enhance: defaults={defaults}')
+ debug_log(f'Prompt enhance: defaults={defaults}')
outputs = self.llm.generate(**inputs, **gen_kwargs)
@@ -855,8 +536,7 @@ class PromptEnhanceScript(scripts_manager.Script):
except Exception as e:
outputs = None
log.error(f'Prompt enhance generate: {e}')
- if debug_enabled:
- errors.display(e, 'Prompt enhance')
+ errors.display(e, 'Prompt enhance')
self.busy = False
response = f'Error: {str(e)}'
finally:
diff --git a/ui/authWrap.ts b/ui/authWrap.ts
index edb16fe57..721946b9a 100644
--- a/ui/authWrap.ts
+++ b/ui/authWrap.ts
@@ -16,7 +16,7 @@ export async function getToken(): Promise<{ user: string | undefined; token: str
const data = (await res.json()) as TokenResponse;
user = data.user;
token = data.token;
- log('getToken', user);
+ log('getToken', { user });
}
}
return { user, token };
diff --git a/ui/css/sdnext.css b/ui/css/sdnext.css
index 1cc088b3b..5975c4771 100644
--- a/ui/css/sdnext.css
+++ b/ui/css/sdnext.css
@@ -2261,8 +2261,12 @@ div:has(>#tab-gallery-folders) {
}
.video-model-link {
- color: var(--button-primary-background-fill);
+ color: var(--button-primary-background-fill) !important;
font-weight: normal;
+ font-size: 0.9em;
+ box-sizing: content-box;
+ position: relative;
+ left: 1em;
}
.controlnet-controls .styler {
diff --git a/ui/dist/sdnext.mjs b/ui/dist/sdnext.mjs
index eee53831e..6dfdfb756 100644
--- a/ui/dist/sdnext.mjs
+++ b/ui/dist/sdnext.mjs
@@ -2759,15 +2759,15 @@ var require_jquery = __commonJS({
function returnFalse() {
return false;
}
- function on(elem, types, selector, data, fn, one) {
+ function on(elem, types2, selector, data, fn, one) {
var origFn, type;
- if (typeof types === "object") {
+ if (typeof types2 === "object") {
if (typeof selector !== "string") {
data = data || selector;
selector = void 0;
}
- for (type in types) {
- on(elem, type, selector, data, types[type], one);
+ for (type in types2) {
+ on(elem, type, selector, data, types2[type], one);
}
return elem;
}
@@ -2798,11 +2798,11 @@ var require_jquery = __commonJS({
fn.guid = origFn.guid || (origFn.guid = jQuery3.guid++);
}
return elem.each(function() {
- jQuery3.event.add(this, types, fn, data, selector);
+ jQuery3.event.add(this, types2, fn, data, selector);
});
}
jQuery3.event = {
- add: function(elem, types, handler, data, selector) {
+ add: function(elem, types2, handler, data, selector) {
var handleObjIn, eventHandle, tmp, events, t, handleObj, special, handlers, type, namespaces, origType, elemData = dataPriv.get(elem);
if (!acceptData(elem)) {
return;
@@ -2826,10 +2826,10 @@ var require_jquery = __commonJS({
return typeof jQuery3 !== "undefined" && jQuery3.event.triggered !== e.type ? jQuery3.event.dispatch.apply(elem, arguments) : void 0;
};
}
- types = (types || "").match(rnothtmlwhite) || [""];
- t = types.length;
+ types2 = (types2 || "").match(rnothtmlwhite) || [""];
+ t = types2.length;
while (t--) {
- tmp = rtypenamespace.exec(types[t]) || [];
+ tmp = rtypenamespace.exec(types2[t]) || [];
type = origType = tmp[1];
namespaces = (tmp[2] || "").split(".").sort();
if (!type) {
@@ -2871,20 +2871,20 @@ var require_jquery = __commonJS({
}
},
// Detach an event or set of events from an element
- remove: function(elem, types, handler, selector, mappedTypes) {
+ remove: function(elem, types2, handler, selector, mappedTypes) {
var j, origCount, tmp, events, t, handleObj, special, handlers, type, namespaces, origType, elemData = dataPriv.hasData(elem) && dataPriv.get(elem);
if (!elemData || !(events = elemData.events)) {
return;
}
- types = (types || "").match(rnothtmlwhite) || [""];
- t = types.length;
+ types2 = (types2 || "").match(rnothtmlwhite) || [""];
+ t = types2.length;
while (t--) {
- tmp = rtypenamespace.exec(types[t]) || [];
+ tmp = rtypenamespace.exec(types2[t]) || [];
type = origType = tmp[1];
namespaces = (tmp[2] || "").split(".").sort();
if (!type) {
for (type in events) {
- jQuery3.event.remove(elem, type + types[t], handler, selector, true);
+ jQuery3.event.remove(elem, type + types2[t], handler, selector, true);
}
continue;
}
@@ -3234,26 +3234,26 @@ var require_jquery = __commonJS({
};
});
jQuery3.fn.extend({
- on: function(types, selector, data, fn) {
- return on(this, types, selector, data, fn);
+ on: function(types2, selector, data, fn) {
+ return on(this, types2, selector, data, fn);
},
- one: function(types, selector, data, fn) {
- return on(this, types, selector, data, fn, 1);
+ one: function(types2, selector, data, fn) {
+ return on(this, types2, selector, data, fn, 1);
},
- off: function(types, selector, fn) {
+ off: function(types2, selector, fn) {
var handleObj, type;
- if (types && types.preventDefault && types.handleObj) {
- handleObj = types.handleObj;
- jQuery3(types.delegateTarget).off(
+ if (types2 && types2.preventDefault && types2.handleObj) {
+ handleObj = types2.handleObj;
+ jQuery3(types2.delegateTarget).off(
handleObj.namespace ? handleObj.origType + "." + handleObj.namespace : handleObj.origType,
handleObj.selector,
handleObj.handler
);
return this;
}
- if (typeof types === "object") {
- for (type in types) {
- this.off(type, selector, types[type]);
+ if (typeof types2 === "object") {
+ for (type in types2) {
+ this.off(type, selector, types2[type]);
}
return this;
}
@@ -3265,7 +3265,7 @@ var require_jquery = __commonJS({
fn = returnFalse;
}
return this.each(function() {
- jQuery3.event.remove(this, types, fn, selector);
+ jQuery3.event.remove(this, types2, fn, selector);
});
}
});
@@ -5823,17 +5823,17 @@ var require_jquery = __commonJS({
};
});
jQuery3.fn.extend({
- bind: function(types, data, fn) {
- return this.on(types, null, data, fn);
+ bind: function(types2, data, fn) {
+ return this.on(types2, null, data, fn);
},
- unbind: function(types, fn) {
- return this.off(types, null, fn);
+ unbind: function(types2, fn) {
+ return this.off(types2, null, fn);
},
- delegate: function(selector, types, data, fn) {
- return this.on(types, selector, data, fn);
+ delegate: function(selector, types2, data, fn) {
+ return this.on(types2, selector, data, fn);
},
- undelegate: function(selector, types, fn) {
- return arguments.length === 1 ? this.off(selector, "**") : this.off(types, selector || "**", fn);
+ undelegate: function(selector, types2, fn) {
+ return arguments.length === 1 ? this.off(selector, "**") : this.off(types2, selector || "**", fn);
},
hover: function(fnOver, fnOut) {
return this.on("mouseenter", fnOver).on("mouseleave", fnOut || fnOver);
@@ -9891,7 +9891,7 @@ async function getToken() {
const data = await res.json();
user = data.user;
token = data.token;
- log("getToken", user);
+ log("getToken", { user });
}
}
return { user, token };
@@ -10025,7 +10025,7 @@ function executeCallbacks(queue, arg) {
const t0 = performance.now();
callback(arg);
const t1 = performance.now();
- if (t1 - t0 > 250) log("callbackSlow", callback.name || callback, `time=${Math.round(t1 - t0)}`);
+ if (t1 - t0 > 250) log("callbackSlow", { callback: callback.name || callback, time: Math.round(t1 - t0) });
timer(callback.name || "anonymousCallback", t1 - t0);
} catch (e) {
error(`executeCallbacks: ${callback} ${e}`);
@@ -10553,12 +10553,12 @@ function sortExtraNetworks(fixed = "no") {
}
const desc = sortDesc[sortVal];
const t1 = performance.now();
- log("sortNetworks", { name: pagename, val: sortVal, order: desc, fixed: fixed === "fixed", items: num, time: Math.round(t1 - t0) });
+ log("sortNetworks", { page: pagename, key: sortVal, order: desc, items: num, time: Math.round(t1 - t0) });
timer(`sortExtraNetworks:${desc}`, t1 - t0);
return desc;
}
async function markSelectedCards(selected, page = "") {
- log("markSelectedCards", selected, page);
+ log("markSelectedCards", { page, selected });
selectedNetworks[page] = selected;
gradioApp().querySelectorAll(".extra-network-cards .card").forEach((el2) => {
if (page.length > 0 && el2.dataset.page !== page) return;
@@ -10578,7 +10578,7 @@ function extractLoraNames(prompt) {
}
function cardClicked(textToAdd) {
const tabName = getENActiveTab();
- log("cardClicked", tabName, textToAdd);
+ log("cardClicked", { tab: tabName, text: textToAdd });
const textarea = activePromptTextarea[tabName];
if (textarea.value.indexOf(textToAdd) !== -1) textarea.value = textarea.value.replace(textToAdd, "");
else textarea.value += textToAdd;
@@ -10993,6 +10993,14 @@ function setRefreshInterval() {
else refreshInterval = window.opts.live_preview_refresh_period || 1e3;
});
}
+function pad2(x) {
+ return x < 10 ? `0${x}` : x;
+}
+function formatTime(secs) {
+ if (secs > 3600) return `${pad2(Math.floor(secs / 60 / 60))}:${pad2(Math.floor(secs / 60) % 60)}:${pad2(Math.floor(secs) % 60)}`;
+ if (secs > 60) return `${pad2(Math.floor(secs / 60))}:${pad2(Math.floor(secs) % 60)}`;
+ return `${Math.floor(secs)}s`;
+}
function checkPaused(state) {
lastState.paused = state ? !state : !lastState.paused;
const t_el = document.getElementById("txt2img_pause");
@@ -11023,21 +11031,50 @@ function setProgress(res) {
eta = min > 0 ? `${Math.round(min)}m ${Math.round(sec)}s` : `${Math.round(sec)}s`;
}
}
+ const elPerf = document.getElementById("control-performance");
+ let hint = "";
+ if (elPerf && res) {
+ const jobTxt = res.job && res.job !== "" ? ` | Job ${res.job}` : "";
+ const batchTxt = res.batch > 0 ? ` | Batch ${res.batch}/${res.batches}` : "";
+ const stateTxt = res.queued ? "Queued" : res.paused ? "Paused" : res.completed ? "Completed" : res.active ? "Active" : "Idle";
+ const stepsTxt = res.step > 0 ? ` | Step ${res.step}/${res.steps}` : "";
+ const progressTxt = res.progress > 0 ? ` | Progress ${Math.round(100 * res.progress)}%` : "";
+ const etaTxt = res.eta > 0 ? ` | ETA ${formatTime(res.eta)}` : "";
+ const previewTxt = res.id_live_preview > 0 ? ` | Preview ${res.id_live_preview}` : "";
+ const elapsedTxt = res.job_time > 0 ? ` | Elapsed ${formatTime(Date.now() / 1e3 - res.job_time)}` : "";
+ const startedTxt = res.job_time > 0 ? ` | Started ${new Date(res.job_time * 1e3).toLocaleTimeString()}` : "";
+ hint = `\u23F1 State ${stateTxt} ${jobTxt} ${startedTxt} ${elapsedTxt} ${batchTxt} ${progressTxt} ${stepsTxt} ${etaTxt} ${previewTxt}`.replaceAll(" ", " ").trim();
+ elPerf.innerHTML = `${hint}`;
+ }
document.title = `SD.Next ${perc}`;
for (const elId of elements) {
- const el2 = document.getElementById(elId);
- if (!el2) continue;
+ const el3 = document.getElementById(elId);
+ if (!el3) continue;
const jobLabel = (res ? `${job} ${perc}${eta}` : "Generate").trim();
- el2.innerText = jobLabel;
+ el3.innerText = jobLabel;
+ el3.title = hint.length > 0 ? hint : jobLabel;
if (!window.waitForUiReady) {
const gradient = perc !== "" ? perc : "100%";
- if (jobLabel === "Generate") el2.style.background = "var(--primary-500)";
+ if (jobLabel === "Generate") el3.style.background = "var(--primary-500)";
else if (jobLabel.endsWith("Decode")) continue;
- else if (jobLabel.endsWith("Start") || jobLabel.endsWith("Finishing")) el2.style.background = "var(--primary-800)";
- else if (res && progress > 0 && progress < 1) el2.style.background = `linear-gradient(to right, var(--primary-500) 0%, var(--primary-800) ${gradient}, var(--neutral-700) ${gradient})`;
- else el2.style.background = "var(--primary-500)";
+ else if (jobLabel.endsWith("Start") || jobLabel.endsWith("Finishing")) el3.style.background = "var(--primary-800)";
+ else if (res && progress > 0 && progress < 1) el3.style.background = `linear-gradient(to right, var(--primary-500) 0%, var(--primary-800) ${gradient}, var(--neutral-700) ${gradient})`;
+ else el3.style.background = "var(--primary-500)";
}
}
+ const el2 = document.getElementById("control-performance");
+ if (el2 && res) {
+ const jobTxt = res.job && res.job !== "" ? ` | Job ${res.job}` : "";
+ const batchTxt = res.batch > 0 ? ` | Batch ${res.batch}/${res.batches}` : "";
+ const stateTxt = res.queued ? "Queued" : res.paused ? "Paused" : res.completed ? "Completed" : res.active ? "Active" : "Idle";
+ const stepsTxt = res.step > 0 ? ` | Step ${res.step}/${res.steps}` : "";
+ const progressTxt = res.progress > 0 ? ` | Progress ${Math.round(100 * res.progress)}%` : "";
+ const etaTxt = res.eta > 0 ? ` | ETA ${formatTime(res.eta)}` : "";
+ const previewTxt = res.id_live_preview > 0 ? ` | Preview ${res.id_live_preview}` : "";
+ const elapsedTxt = res.job_time > 0 ? ` | Elapsed ${formatTime(Date.now() / 1e3 - res.job_time)}` : "";
+ const startedTxt = res.job_time > 0 ? ` | Started ${new Date(res.job_time * 1e3).toLocaleTimeString()}` : "";
+ el2.innerHTML = `
\u23F1 State ${stateTxt} ${jobTxt} ${startedTxt} ${elapsedTxt} ${batchTxt} ${progressTxt} ${stepsTxt} ${etaTxt} ${previewTxt}
`.replaceAll(" ", " ").trim();
+ }
}
function requestInterrupt() {
setProgress();
@@ -11109,56 +11146,65 @@ function requestProgress(id_task = "undefined", progressEl = null, galleryEl = n
return true;
}
};
+ const onProgressDataHandler = async (res, caller) => {
+ if (res?.debug) debug("progress:", { start: dateStart, res });
+ lastState = res;
+ const elapsedFromStart = (Date.now() - dateStart) / 1e3;
+ hasStarted = hasStarted || res.active;
+ if (res.completed || !res.active && (hasStarted || once)) {
+ debug("progress", { end: res, reason: res.completed ? "completed" : "inactive" });
+ const hidden = document.hidden || !previewVisible();
+ if (!res.paused) removeLivePreview(!hidden);
+ return;
+ }
+ if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) {
+ debug("progress", { end: res, reason: "progressTimeout" });
+ if (!res.paused) removeLivePreview(false);
+ return;
+ }
+ if (elapsedFromStart > startTimeout && !res.queued && !res.active) {
+ debug("progress", { end: res, reason: "startTimeout" });
+ if (!res.paused) removeLivePreview(false);
+ return;
+ }
+ if (res.progress !== prevProgress) {
+ dateStart = Date.now();
+ prevProgress = res.progress;
+ }
+ setProgress(res);
+ if (res.live_preview && !livePreview) initLivePreview();
+ let id_live_preview = res.id_live_preview;
+ if (res.live_preview && galleryEl) {
+ if (img.src !== res.live_preview) img.src = res.live_preview;
+ id_live_preview = res.id_live_preview;
+ }
+ if (onProgress) onProgress(res);
+ let timeout = Math.max(window.opts.live_preview_refresh_period || 500, 500);
+ timeout += (Math.random() * 0.4 - 0.2) * timeout;
+ setTimeout(() => caller(id_task, id_live_preview), timeout);
+ };
+ const onProgressErrorHandler = (err) => {
+ error("progress", { error: err });
+ removeLivePreview(false);
+ };
const startLivePreview = (taskId, id_live_preview) => {
- if (window.opts.live_preview_refresh_period === 0) return;
- let request_id = -1;
const hidden = document.hidden || !previewVisible();
+ let request_id = id_live_preview;
if (hidden) {
if (!window.opts.live_preview_require_focus) request_id = id_live_preview;
- } else {
- request_id = id_live_preview;
+ } else if (window.opts.live_preview_refresh_period === 0) {
+ request_id = -1;
}
- const onProgressHandler = (res) => {
- if (res?.debug) debug("progress:", { start: dateStart, id: request_id, res });
- lastState = res;
- const elapsedFromStart = (Date.now() - dateStart) / 1e3;
- hasStarted = hasStarted || res.active;
- if (res.completed || !res.active && (hasStarted || once)) {
- debug("progress", { end: res, reason: res.completed ? "completed" : "inactive" });
- if (!res.paused) removeLivePreview(!hidden);
- return;
- }
- if (elapsedFromStart > progressTimeout && !res.queued && res.progress === prevProgress) {
- debug("progress", { end: res, reason: "progressTimeout" });
- if (!res.paused) removeLivePreview(false);
- return;
- }
- if (elapsedFromStart > startTimeout && !res.queued && !res.active) {
- debug("progress", { end: res, reason: "startTimeout" });
- if (!res.paused) removeLivePreview(false);
- return;
- }
- if (res.progress !== prevProgress) {
- dateStart = Date.now();
- prevProgress = res.progress;
- }
- setProgress(res);
- if (res.live_preview && !livePreview) initLivePreview();
- if (res.live_preview && galleryEl) {
- if (img.src !== res.live_preview) img.src = res.live_preview;
- id_live_preview = res.id_live_preview;
- }
- if (onProgress) onProgress(res);
- setTimeout(() => startLivePreview(id_task, id_live_preview), window.opts.live_preview_refresh_period || 500);
- };
- const onProgressErrorHandler = (err) => {
- error("progress", { error: err });
- removeLivePreview(false);
- };
- xhrPost("./internal/progress", { id_task, id_live_preview: request_id }, onProgressHandler, onProgressErrorHandler, false, 3e4);
+ xhrPost("./internal/progress", { id_task, id_live_preview: request_id }, onLivePreviewHandler, onProgressErrorHandler, false, 3e4);
};
+ const startProgress = (taskId, id_live_preview) => {
+ xhrPost("./internal/progress", { id_task, id_live_preview: -1 }, onProgressHandler, onProgressErrorHandler, false, 3e4);
+ };
+ const onProgressHandler = (res) => onProgressDataHandler(res, startProgress);
+ const onLivePreviewHandler = (res) => onProgressDataHandler(res, startLivePreview);
debug("progress", { start: dateStart });
startLivePreview(id_task, 0);
+ startProgress(id_task, -1);
}
window.checkPaused = checkPaused;
window.requestInterrupt = requestInterrupt;
@@ -12424,10 +12470,11 @@ async function initSettings() {
// ui/monitor.ts
var monitorActive = false;
+var wsTimer;
var ConnectionMonitorState = class _ConnectionMonitorState {
static ws;
static url = "";
- static delay = 1e3;
+ static delay = 2e3;
static element;
static version = "";
static commit = "";
@@ -12447,7 +12494,7 @@ var ConnectionMonitorState = class _ConnectionMonitorState {
if (online !== this.online) {
this.online = online;
this.ts = /* @__PURE__ */ new Date();
- debug("monitorState", { online: _ConnectionMonitorState.online, ts: _ConnectionMonitorState.ts });
+ debug("monitorState", { online: _ConnectionMonitorState.online, ts: _ConnectionMonitorState.ts?.toLocaleTimeString() });
}
if (data?.updated) this.version = data.updated;
if (data?.commit) this.commit = data.commit;
@@ -12480,21 +12527,41 @@ async function updateIndicator(online, data = {}, msg) {
ConnectionMonitorState.updateState();
if (msg) log("monitorConnection:", { online, data, msg });
}
+function scheduleNextLoop() {
+ if (wsTimer) {
+ clearTimeout(wsTimer);
+ wsTimer = void 0;
+ }
+ const offlineDurationMs = Date.now() - ConnectionMonitorState.ts.getTime();
+ if (!ConnectionMonitorState.online && offlineDurationMs > 60 * 60 * 1e3) ConnectionMonitorState.delay = 1e4;
+ else if (!ConnectionMonitorState.online && offlineDurationMs > 5 * 60 * 1e3) ConnectionMonitorState.delay = 5e3;
+ else ConnectionMonitorState.delay = 1e3;
+ wsTimer = setTimeout(wsMonitorLoop, ConnectionMonitorState.delay);
+}
async function wsMonitorLoop() {
- const delayed = Date.now() - ConnectionMonitorState.ts.getTime();
- if (delayed > 60 * 60 && ConnectionMonitorState.delay < 10 && !ConnectionMonitorState.online) ConnectionMonitorState.delay = 1e4;
- else if (delayed > 5 * 60 && ConnectionMonitorState.delay < 5 && !ConnectionMonitorState.online) ConnectionMonitorState.delay = 5e3;
- else ConnectionMonitorState.delay = 2e3;
+ if (ConnectionMonitorState.ws) {
+ ConnectionMonitorState.ws.onopen = null;
+ ConnectionMonitorState.ws.onmessage = null;
+ ConnectionMonitorState.ws.onclose = null;
+ ConnectionMonitorState.ws.onerror = null;
+ try {
+ ConnectionMonitorState.ws.close();
+ } catch {
+ }
+ ConnectionMonitorState.ws = void 0;
+ }
try {
- ConnectionMonitorState.ws = new WebSocket(`${ConnectionMonitorState.url}/queue/join`);
- ConnectionMonitorState.ws.onopen = () => {
+ ConnectionMonitorState.ws = new WebSocket(`${ConnectionMonitorState.url}/internal/monitor`);
+ ConnectionMonitorState.ws.onopen = () => updateIndicator(true);
+ ConnectionMonitorState.ws.onmessage = (msg) => updateIndicator(true, msg.data ? JSON.parse(msg.data) : {});
+ ConnectionMonitorState.ws.onclose = () => {
+ updateIndicator(false);
+ scheduleNextLoop();
};
- ConnectionMonitorState.ws.onmessage = () => updateIndicator(true);
- ConnectionMonitorState.ws.onclose = () => setTimeout(wsMonitorLoop, ConnectionMonitorState.delay);
ConnectionMonitorState.ws.onerror = (e) => updateIndicator(false, {}, String(e.message || "unknown error"));
} catch (e) {
updateIndicator(false, {}, String(e.message || e));
- setTimeout(monitorConnection, ConnectionMonitorState.delay);
+ scheduleNextLoop();
}
}
async function monitorConnection() {
@@ -12518,7 +12585,7 @@ async function monitorConnection() {
wsMonitorLoop();
} catch {
updateIndicator(false, data);
- setTimeout(monitorConnection, ConnectionMonitorState.delay);
+ scheduleNextLoop();
}
}
@@ -16448,7 +16515,7 @@ var Timesheet = class {
// ui/history.ts
var inferenceTypes = ["inference", "vae", "te"];
var ioTypes = ["load", "save"];
-function refreshHistory() {
+async function refreshHistory() {
log("refreshHistory");
authFetch(`${window.api}/history`, { priority: "low" }).then((res) => {
if (!res) return;
@@ -16500,6 +16567,70 @@ function refreshHistory() {
}
window.refreshHistory = refreshHistory;
+// ui/storage.ts
+var types = ["Images", "Videos", "Models", "Data", "Cache", "Code", "Other"];
+function buildTable(type, data) {
+ const totalSize = data.reduce((acc, entry) => acc + entry.size, 0);
+ const totalLoc = data.length;
+ const totalFiles = data.reduce((acc, entry) => acc + entry.nfiles, 0);
+ const totalFolders = data.reduce((acc, entry) => acc + entry.nfolders, 0);
+ let title = `Locations: ${totalLoc}
+Total Size: ${(totalSize / (1024 * 1024)).toLocaleString(void 0, { minimumFractionDigits: 2, maximumFractionDigits: 2 })} MB
+Total Files: ${totalFiles}
+Total Folders: ${totalFolders}
+`;
+ let html = `${type}
`;
+ for (const entry of data) {
+ if (entry.size === 0) continue;
+ const size = (entry.size / (1024 * 1024)).toLocaleString(void 0, { minimumFractionDigits: 2, maximumFractionDigits: 2 }) + " MB";
+ const mtime = entry.mtime > 0 ? new Date(entry.mtime * 1e3).toLocaleString() : "";
+ title = `Type: ${entry.type}
+Name: ${entry.name}
+Size: ${size}
+Last modified: ${mtime}
+`;
+ title += `Folders: ${entry.folders.join(", ")}
+Resolved paths: ${entry.paths.join(", ")}
+`;
+ title += `Subfolders: ${entry.nfolders}
+Files: ${entry.nfiles}
+Symlinks: ${entry.nsymlinks}
+Errors: ${entry.nerrors}
+`;
+ title += `Time to scan: ${entry.time.toFixed(3)} seconds`;
+ const perc = Math.round(entry.size / totalSize * 100);
+ const color = `rgb(${perc}, 50, 80)`;
+ const css = `background: linear-gradient(to right, ${color} ${perc}%, transparent ${perc}%);`;
+ html += `| ${entry.name} | ${size} | ${mtime} |
`;
+ }
+ html += "
";
+ return html;
+}
+async function refreshStorage(storageTypes) {
+ log("refreshStorage", storageTypes);
+ authFetch(`${window.api}/storage?types=${storageTypes.join(",")}`, { priority: "low" }).then((res) => {
+ if (!res) return;
+ const timeline = document.getElementById("storage_timeline");
+ const table = document.getElementById("storage_table");
+ if (!timeline || !table) return;
+ timeline.innerHTML = "";
+ res.json().then((rawData) => {
+ const data = rawData;
+ if (!data || !data.length) {
+ table.innerHTML = "No storage data available.
";
+ return;
+ }
+ table.innerHTML = "";
+ if (storageTypes.includes("All")) storageTypes = types;
+ for (const type of storageTypes) {
+ const typeData = data.filter((entry) => entry.type === type);
+ if (typeData.length > 0) table.innerHTML += buildTable(type, typeData);
+ }
+ });
+ });
+}
+window.refreshStorage = refreshStorage;
+
// ui/aspectRatioOverlay.ts
var currentWidth = null;
var currentHeight = null;
diff --git a/ui/dist/sdnext.mjs.map b/ui/dist/sdnext.mjs.map
index 8affd322d..daf92b70a 100644
--- a/ui/dist/sdnext.mjs.map
+++ b/ui/dist/sdnext.mjs.map
@@ -1,7 +1,7 @@
{
"version": 3,
- "sources": ["../../node_modules/.pnpm/jquery@4.0.0/node_modules/jquery/dist/jquery.js", "../js/iframeResizer.js", "../../node_modules/.pnpm/exifr@7.1.3/node_modules/exifr/dist/full.umd.js", "../../node_modules/.pnpm/wheel@1.0.0/node_modules/wheel/index.js", "../../node_modules/.pnpm/bezier-easing@2.1.0/node_modules/bezier-easing/src/index.js", "../../node_modules/.pnpm/amator@1.1.0/node_modules/amator/index.js", "../../node_modules/.pnpm/ngraph.events@1.4.0/node_modules/ngraph.events/index.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/lib/kinetic.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/lib/makeTextSelectionInterceptor.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/lib/transform.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/lib/makeSvgController.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/lib/makeDomController.js", "../../node_modules/.pnpm/panzoom@9.4.4/node_modules/panzoom/index.js", "../../node_modules/.pnpm/jquery@4.0.0/node_modules/jquery/dist-module/wrappers/jquery.node-module-wrapper.js", "../vendor.ts", "../logger.ts", "../authWrap.ts", "../timers.ts", "../script.ts", "../changelog.ts", "../control.ts", "../extraNetworks.ts", "../generationParams.ts", "../imageParams.ts", "../notification.ts", "../progressBar.ts", "../ui.ts", "../inputAccordion.ts", "../indexdb.ts", "../logMonitor.ts", "../settings.ts", "../monitor.ts", "../promptChecker.ts", "../js/sha256.ts", "../gallery.ts", "../imageViewer.ts", "../autocomplete_xn.ts", "../autocomplete.ts", "../setHints.ts", "../contextMenus.ts", "../uiConfig.ts", "../loader.ts", "../legacy.ts", "../startup.ts", "../extensions.ts", "../dragDrop.ts", "../civitai.ts", "../guidance.ts", "../timesheet.ts", "../history.ts", "../aspectRatioOverlay.ts", "../resolutionLock.ts", "../editAttention.ts", "../../node_modules/.pnpm/jquery-sparkline@2.4.0/node_modules/jquery-sparkline/jquery.sparkline.js", "../gpu.ts"],
- "sourcesContent": ["/*!\n * jQuery JavaScript Library v4.0.0\n * https://jquery.com/\n *\n * Copyright OpenJS Foundation and other contributors\n * Released under the MIT license\n * https://jquery.com/license/\n *\n * Date: 2026-01-18T00:20Z\n */\n( function( global, factory ) {\n\n\t\"use strict\";\n\n\tif ( typeof module === \"object\" && typeof module.exports === \"object\" ) {\n\n\t\t// For CommonJS and CommonJS-like environments where a proper `window`\n\t\t// is present, execute the factory and get jQuery.\n\t\tmodule.exports = factory( global, true );\n\t} else {\n\t\tfactory( global );\n\t}\n\n// Pass this if window is not defined yet\n} )( typeof window !== \"undefined\" ? window : this, function( window, noGlobal ) {\n\n\"use strict\";\n\nif ( !window.document ) {\n\tthrow new Error( \"jQuery requires a window with a document\" );\n}\n\nvar arr = [];\n\nvar getProto = Object.getPrototypeOf;\n\nvar slice = arr.slice;\n\n// Support: IE 11+\n// IE doesn't have Array#flat; provide a fallback.\nvar flat = arr.flat ? function( array ) {\n\treturn arr.flat.call( array );\n} : function( array ) {\n\treturn arr.concat.apply( [], array );\n};\n\nvar push = arr.push;\n\nvar indexOf = arr.indexOf;\n\n// [[Class]] -> type pairs\nvar class2type = {};\n\nvar toString = class2type.toString;\n\nvar hasOwn = class2type.hasOwnProperty;\n\nvar fnToString = hasOwn.toString;\n\nvar ObjectFunctionString = fnToString.call( Object );\n\n// All support tests are defined in their respective modules.\nvar support = {};\n\nfunction toType( obj ) {\n\tif ( obj == null ) {\n\t\treturn obj + \"\";\n\t}\n\n\treturn typeof obj === \"object\" ?\n\t\tclass2type[ toString.call( obj ) ] || \"object\" :\n\t\ttypeof obj;\n}\n\nfunction isWindow( obj ) {\n\treturn obj != null && obj === obj.window;\n}\n\nfunction isArrayLike( obj ) {\n\n\tvar length = !!obj && obj.length,\n\t\ttype = toType( obj );\n\n\tif ( typeof obj === \"function\" || isWindow( obj ) ) {\n\t\treturn false;\n\t}\n\n\treturn type === \"array\" || length === 0 ||\n\t\ttypeof length === \"number\" && length > 0 && ( length - 1 ) in obj;\n}\n\nvar document$1 = window.document;\n\nvar preservedScriptAttributes = {\n\ttype: true,\n\tsrc: true,\n\tnonce: true,\n\tnoModule: true\n};\n\nfunction DOMEval( code, node, doc ) {\n\tdoc = doc || document$1;\n\n\tvar i,\n\t\tscript = doc.createElement( \"script\" );\n\n\tscript.text = code;\n\tfor ( i in preservedScriptAttributes ) {\n\t\tif ( node && node[ i ] ) {\n\t\t\tscript[ i ] = node[ i ];\n\t\t}\n\t}\n\n\tif ( doc.head.appendChild( script ).parentNode ) {\n\t\tscript.parentNode.removeChild( script );\n\t}\n}\n\nvar version = \"4.0.0\",\n\n\trhtmlSuffix = /HTML$/i,\n\n\t// Define a local copy of jQuery\n\tjQuery = function( selector, context ) {\n\n\t\t// The jQuery object is actually just the init constructor 'enhanced'\n\t\t// Need init if jQuery is called (just allow error to be thrown if not included)\n\t\treturn new jQuery.fn.init( selector, context );\n\t};\n\njQuery.fn = jQuery.prototype = {\n\n\t// The current version of jQuery being used\n\tjquery: version,\n\n\tconstructor: jQuery,\n\n\t// The default length of a jQuery object is 0\n\tlength: 0,\n\n\ttoArray: function() {\n\t\treturn slice.call( this );\n\t},\n\n\t// Get the Nth element in the matched element set OR\n\t// Get the whole matched element set as a clean array\n\tget: function( num ) {\n\n\t\t// Return all the elements in a clean array\n\t\tif ( num == null ) {\n\t\t\treturn slice.call( this );\n\t\t}\n\n\t\t// Return just the one element from the set\n\t\treturn num < 0 ? this[ num + this.length ] : this[ num ];\n\t},\n\n\t// Take an array of elements and push it onto the stack\n\t// (returning the new matched element set)\n\tpushStack: function( elems ) {\n\n\t\t// Build a new jQuery matched element set\n\t\tvar ret = jQuery.merge( this.constructor(), elems );\n\n\t\t// Add the old object onto the stack (as a reference)\n\t\tret.prevObject = this;\n\n\t\t// Return the newly-formed element set\n\t\treturn ret;\n\t},\n\n\t// Execute a callback for every element in the matched set.\n\teach: function( callback ) {\n\t\treturn jQuery.each( this, callback );\n\t},\n\n\tmap: function( callback ) {\n\t\treturn this.pushStack( jQuery.map( this, function( elem, i ) {\n\t\t\treturn callback.call( elem, i, elem );\n\t\t} ) );\n\t},\n\n\tslice: function() {\n\t\treturn this.pushStack( slice.apply( this, arguments ) );\n\t},\n\n\tfirst: function() {\n\t\treturn this.eq( 0 );\n\t},\n\n\tlast: function() {\n\t\treturn this.eq( -1 );\n\t},\n\n\teven: function() {\n\t\treturn this.pushStack( jQuery.grep( this, function( _elem, i ) {\n\t\t\treturn ( i + 1 ) % 2;\n\t\t} ) );\n\t},\n\n\todd: function() {\n\t\treturn this.pushStack( jQuery.grep( this, function( _elem, i ) {\n\t\t\treturn i % 2;\n\t\t} ) );\n\t},\n\n\teq: function( i ) {\n\t\tvar len = this.length,\n\t\t\tj = +i + ( i < 0 ? len : 0 );\n\t\treturn this.pushStack( j >= 0 && j < len ? [ this[ j ] ] : [] );\n\t},\n\n\tend: function() {\n\t\treturn this.prevObject || this.constructor();\n\t}\n};\n\njQuery.extend = jQuery.fn.extend = function() {\n\tvar options, name, src, copy, copyIsArray, clone,\n\t\ttarget = arguments[ 0 ] || {},\n\t\ti = 1,\n\t\tlength = arguments.length,\n\t\tdeep = false;\n\n\t// Handle a deep copy situation\n\tif ( typeof target === \"boolean\" ) {\n\t\tdeep = target;\n\n\t\t// Skip the boolean and the target\n\t\ttarget = arguments[ i ] || {};\n\t\ti++;\n\t}\n\n\t// Handle case when target is a string or something (possible in deep copy)\n\tif ( typeof target !== \"object\" && typeof target !== \"function\" ) {\n\t\ttarget = {};\n\t}\n\n\t// Extend jQuery itself if only one argument is passed\n\tif ( i === length ) {\n\t\ttarget = this;\n\t\ti--;\n\t}\n\n\tfor ( ; i < length; i++ ) {\n\n\t\t// Only deal with non-null/undefined values\n\t\tif ( ( options = arguments[ i ] ) != null ) {\n\n\t\t\t// Extend the base object\n\t\t\tfor ( name in options ) {\n\t\t\t\tcopy = options[ name ];\n\n\t\t\t\t// Prevent Object.prototype pollution\n\t\t\t\t// Prevent never-ending loop\n\t\t\t\tif ( name === \"__proto__\" || target === copy ) {\n\t\t\t\t\tcontinue;\n\t\t\t\t}\n\n\t\t\t\t// Recurse if we're merging plain objects or arrays\n\t\t\t\tif ( deep && copy && ( jQuery.isPlainObject( copy ) ||\n\t\t\t\t\t( copyIsArray = Array.isArray( copy ) ) ) ) {\n\t\t\t\t\tsrc = target[ name ];\n\n\t\t\t\t\t// Ensure proper type for the source value\n\t\t\t\t\tif ( copyIsArray && !Array.isArray( src ) ) {\n\t\t\t\t\t\tclone = [];\n\t\t\t\t\t} else if ( !copyIsArray && !jQuery.isPlainObject( src ) ) {\n\t\t\t\t\t\tclone = {};\n\t\t\t\t\t} else {\n\t\t\t\t\t\tclone = src;\n\t\t\t\t\t}\n\t\t\t\t\tcopyIsArray = false;\n\n\t\t\t\t\t// Never move original objects, clone them\n\t\t\t\t\ttarget[ name ] = jQuery.extend( deep, clone, copy );\n\n\t\t\t\t// Don't bring in undefined values\n\t\t\t\t} else if ( copy !== undefined ) {\n\t\t\t\t\ttarget[ name ] = copy;\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\t// Return the modified object\n\treturn target;\n};\n\njQuery.extend( {\n\n\t// Unique for each copy of jQuery on the page\n\texpando: \"jQuery\" + ( version + Math.random() ).replace( /\\D/g, \"\" ),\n\n\t// Assume jQuery is ready without the ready module\n\tisReady: true,\n\n\terror: function( msg ) {\n\t\tthrow new Error( msg );\n\t},\n\n\tnoop: function() {},\n\n\tisPlainObject: function( obj ) {\n\t\tvar proto, Ctor;\n\n\t\t// Detect obvious negatives\n\t\t// Use toString instead of jQuery.type to catch host objects\n\t\tif ( !obj || toString.call( obj ) !== \"[object Object]\" ) {\n\t\t\treturn false;\n\t\t}\n\n\t\tproto = getProto( obj );\n\n\t\t// Objects with no prototype (e.g., `Object.create( null )`) are plain\n\t\tif ( !proto ) {\n\t\t\treturn true;\n\t\t}\n\n\t\t// Objects with prototype are plain iff they were constructed by a global Object function\n\t\tCtor = hasOwn.call( proto, \"constructor\" ) && proto.constructor;\n\t\treturn typeof Ctor === \"function\" && fnToString.call( Ctor ) === ObjectFunctionString;\n\t},\n\n\tisEmptyObject: function( obj ) {\n\t\tvar name;\n\n\t\tfor ( name in obj ) {\n\t\t\treturn false;\n\t\t}\n\t\treturn true;\n\t},\n\n\t// Evaluates a script in a provided context; falls back to the global one\n\t// if not specified.\n\tglobalEval: function( code, options, doc ) {\n\t\tDOMEval( code, { nonce: options && options.nonce }, doc );\n\t},\n\n\teach: function( obj, callback ) {\n\t\tvar length, i = 0;\n\n\t\tif ( isArrayLike( obj ) ) {\n\t\t\tlength = obj.length;\n\t\t\tfor ( ; i < length; i++ ) {\n\t\t\t\tif ( callback.call( obj[ i ], i, obj[ i ] ) === false ) {\n\t\t\t\t\tbreak;\n\t\t\t\t}\n\t\t\t}\n\t\t} else {\n\t\t\tfor ( i in obj ) {\n\t\t\t\tif ( callback.call( obj[ i ], i, obj[ i ] ) === false ) {\n\t\t\t\t\tbreak;\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\treturn obj;\n\t},\n\n\n\t// Retrieve the text value of an array of DOM nodes\n\ttext: function( elem ) {\n\t\tvar node,\n\t\t\tret = \"\",\n\t\t\ti = 0,\n\t\t\tnodeType = elem.nodeType;\n\n\t\tif ( !nodeType ) {\n\n\t\t\t// If no nodeType, this is expected to be an array\n\t\t\twhile ( ( node = elem[ i++ ] ) ) {\n\n\t\t\t\t// Do not traverse comment nodes\n\t\t\t\tret += jQuery.text( node );\n\t\t\t}\n\t\t}\n\t\tif ( nodeType === 1 || nodeType === 11 ) {\n\t\t\treturn elem.textContent;\n\t\t}\n\t\tif ( nodeType === 9 ) {\n\t\t\treturn elem.documentElement.textContent;\n\t\t}\n\t\tif ( nodeType === 3 || nodeType === 4 ) {\n\t\t\treturn elem.nodeValue;\n\t\t}\n\n\t\t// Do not include comment or processing instruction nodes\n\n\t\treturn ret;\n\t},\n\n\n\t// results is for internal usage only\n\tmakeArray: function( arr, results ) {\n\t\tvar ret = results || [];\n\n\t\tif ( arr != null ) {\n\t\t\tif ( isArrayLike( Object( arr ) ) ) {\n\t\t\t\tjQuery.merge( ret,\n\t\t\t\t\ttypeof arr === \"string\" ?\n\t\t\t\t\t\t[ arr ] : arr\n\t\t\t\t);\n\t\t\t} else {\n\t\t\t\tpush.call( ret, arr );\n\t\t\t}\n\t\t}\n\n\t\treturn ret;\n\t},\n\n\tinArray: function( elem, arr, i ) {\n\t\treturn arr == null ? -1 : indexOf.call( arr, elem, i );\n\t},\n\n\tisXMLDoc: function( elem ) {\n\t\tvar namespace = elem && elem.namespaceURI,\n\t\t\tdocElem = elem && ( elem.ownerDocument || elem ).documentElement;\n\n\t\t// Assume HTML when documentElement doesn't yet exist, such as inside\n\t\t// document fragments.\n\t\treturn !rhtmlSuffix.test( namespace || docElem && docElem.nodeName || \"HTML\" );\n\t},\n\n\t// Note: an element does not contain itself\n\tcontains: function( a, b ) {\n\t\tvar bup = b && b.parentNode;\n\n\t\treturn a === bup || !!( bup && bup.nodeType === 1 && (\n\n\t\t\t// Support: IE 9 - 11+\n\t\t\t// IE doesn't have `contains` on SVG.\n\t\t\ta.contains ?\n\t\t\t\ta.contains( bup ) :\n\t\t\t\ta.compareDocumentPosition && a.compareDocumentPosition( bup ) & 16\n\t\t) );\n\t},\n\n\tmerge: function( first, second ) {\n\t\tvar len = +second.length,\n\t\t\tj = 0,\n\t\t\ti = first.length;\n\n\t\tfor ( ; j < len; j++ ) {\n\t\t\tfirst[ i++ ] = second[ j ];\n\t\t}\n\n\t\tfirst.length = i;\n\n\t\treturn first;\n\t},\n\n\tgrep: function( elems, callback, invert ) {\n\t\tvar callbackInverse,\n\t\t\tmatches = [],\n\t\t\ti = 0,\n\t\t\tlength = elems.length,\n\t\t\tcallbackExpect = !invert;\n\n\t\t// Go through the array, only saving the items\n\t\t// that pass the validator function\n\t\tfor ( ; i < length; i++ ) {\n\t\t\tcallbackInverse = !callback( elems[ i ], i );\n\t\t\tif ( callbackInverse !== callbackExpect ) {\n\t\t\t\tmatches.push( elems[ i ] );\n\t\t\t}\n\t\t}\n\n\t\treturn matches;\n\t},\n\n\t// arg is for internal usage only\n\tmap: function( elems, callback, arg ) {\n\t\tvar length, value,\n\t\t\ti = 0,\n\t\t\tret = [];\n\n\t\t// Go through the array, translating each of the items to their new values\n\t\tif ( isArrayLike( elems ) ) {\n\t\t\tlength = elems.length;\n\t\t\tfor ( ; i < length; i++ ) {\n\t\t\t\tvalue = callback( elems[ i ], i, arg );\n\n\t\t\t\tif ( value != null ) {\n\t\t\t\t\tret.push( value );\n\t\t\t\t}\n\t\t\t}\n\n\t\t// Go through every key on the object,\n\t\t} else {\n\t\t\tfor ( i in elems ) {\n\t\t\t\tvalue = callback( elems[ i ], i, arg );\n\n\t\t\t\tif ( value != null ) {\n\t\t\t\t\tret.push( value );\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\t// Flatten any nested arrays\n\t\treturn flat( ret );\n\t},\n\n\t// A global GUID counter for objects\n\tguid: 1,\n\n\t// jQuery.support is not used in Core but other projects attach their\n\t// properties to it so it needs to exist.\n\tsupport: support\n} );\n\nif ( typeof Symbol === \"function\" ) {\n\tjQuery.fn[ Symbol.iterator ] = arr[ Symbol.iterator ];\n}\n\n// Populate the class2type map\njQuery.each( \"Boolean Number String Function Array Date RegExp Object Error Symbol\".split( \" \" ),\n\tfunction( _i, name ) {\n\t\tclass2type[ \"[object \" + name + \"]\" ] = name.toLowerCase();\n\t} );\n\nfunction nodeName( elem, name ) {\n\treturn elem.nodeName && elem.nodeName.toLowerCase() === name.toLowerCase();\n}\n\nvar pop = arr.pop;\n\n// https://www.w3.org/TR/css3-selectors/#whitespace\nvar whitespace = \"[\\\\x20\\\\t\\\\r\\\\n\\\\f]\";\n\nvar isIE = document$1.documentMode;\n\nvar rbuggyQSA = isIE && new RegExp(\n\n\t// Support: IE 9 - 11+\n\t// IE's :disabled selector does not pick up the children of disabled fieldsets\n\t\":enabled|:disabled|\" +\n\n\t// Support: IE 11+\n\t// IE 11 doesn't find elements on a `[name='']` query in some cases.\n\t// Adding a temporary attribute to the document before the selection works\n\t// around the issue.\n\t\"\\\\[\" + whitespace + \"*name\" + whitespace + \"*=\" +\n\twhitespace + \"*(?:''|\\\"\\\")\"\n\n);\n\nvar rtrimCSS = new RegExp(\n\t\"^\" + whitespace + \"+|((?:^|[^\\\\\\\\])(?:\\\\\\\\.)*)\" + whitespace + \"+$\",\n\t\"g\"\n);\n\n// https://www.w3.org/TR/css-syntax-3/#ident-token-diagram\nvar identifier = \"(?:\\\\\\\\[\\\\da-fA-F]{1,6}\" + whitespace +\n\t\"?|\\\\\\\\[^\\\\r\\\\n\\\\f]|[\\\\w-]|[^\\0-\\\\x7f])+\";\n\nvar rleadingCombinator = new RegExp( \"^\" + whitespace + \"*([>+~]|\" +\n\twhitespace + \")\" + whitespace + \"*\" );\n\nvar rdescend = new RegExp( whitespace + \"|>\" );\n\nvar rsibling = /[+~]/;\n\nvar documentElement$1 = document$1.documentElement;\n\n// Support: IE 9 - 11+\n// IE requires a prefix.\nvar matches = documentElement$1.matches || documentElement$1.msMatchesSelector;\n\n/**\n * Create key-value caches of limited size\n * @returns {function(string, object)} Returns the Object data after storing it on itself with\n *\tproperty name the (space-suffixed) string and (if the cache is larger than Expr.cacheLength)\n *\tdeleting the oldest entry\n */\nfunction createCache() {\n\tvar keys = [];\n\n\tfunction cache( key, value ) {\n\n\t\t// Use (key + \" \") to avoid collision with native prototype properties\n\t\t// (see https://github.com/jquery/sizzle/issues/157)\n\t\tif ( keys.push( key + \" \" ) > jQuery.expr.cacheLength ) {\n\n\t\t\t// Only keep the most recent entries\n\t\t\tdelete cache[ keys.shift() ];\n\t\t}\n\t\treturn ( cache[ key + \" \" ] = value );\n\t}\n\treturn cache;\n}\n\n/**\n * Checks a node for validity as a jQuery selector context\n * @param {Element|Object=} context\n * @returns {Element|Object|Boolean} The input node if acceptable, otherwise a falsy value\n */\nfunction testContext( context ) {\n\treturn context && typeof context.getElementsByTagName !== \"undefined\" && context;\n}\n\n// Attribute selectors: https://www.w3.org/TR/selectors/#attribute-selectors\nvar attributes = \"\\\\[\" + whitespace + \"*(\" + identifier + \")(?:\" + whitespace +\n\n\t// Operator (capture 2)\n\t\"*([*^$|!~]?=)\" + whitespace +\n\n\t// \"Attribute values must be CSS identifiers [capture 5] or strings [capture 3 or capture 4]\"\n\t\"*(?:'((?:\\\\\\\\.|[^\\\\\\\\'])*)'|\\\"((?:\\\\\\\\.|[^\\\\\\\\\\\"])*)\\\"|(\" + identifier + \"))|)\" +\n\twhitespace + \"*\\\\]\";\n\nvar pseudos = \":(\" + identifier + \")(?:\\\\((\" +\n\n\t// To reduce the number of selectors needing tokenize in the preFilter, prefer arguments:\n\t// 1. quoted (capture 3; capture 4 or capture 5)\n\t\"('((?:\\\\\\\\.|[^\\\\\\\\'])*)'|\\\"((?:\\\\\\\\.|[^\\\\\\\\\\\"])*)\\\")|\" +\n\n\t// 2. simple (capture 6)\n\t\"((?:\\\\\\\\.|[^\\\\\\\\()[\\\\]]|\" + attributes + \")*)|\" +\n\n\t// 3. anything else (capture 2)\n\t\".*\" +\n\t\")\\\\)|)\";\n\nvar filterMatchExpr = {\n\tID: new RegExp( \"^#(\" + identifier + \")\" ),\n\tCLASS: new RegExp( \"^\\\\.(\" + identifier + \")\" ),\n\tTAG: new RegExp( \"^(\" + identifier + \"|[*])\" ),\n\tATTR: new RegExp( \"^\" + attributes ),\n\tPSEUDO: new RegExp( \"^\" + pseudos ),\n\tCHILD: new RegExp(\n\t\t\"^:(only|first|last|nth|nth-last)-(child|of-type)(?:\\\\(\" +\n\t\twhitespace + \"*(even|odd|(([+-]|)(\\\\d*)n|)\" + whitespace + \"*(?:([+-]|)\" +\n\t\twhitespace + \"*(\\\\d+)|))\" + whitespace + \"*\\\\)|)\", \"i\" )\n};\n\nvar rpseudo = new RegExp( pseudos );\n\n// CSS escapes\n// https://www.w3.org/TR/CSS21/syndata.html#escaped-characters\n\nvar runescape = new RegExp( \"\\\\\\\\[\\\\da-fA-F]{1,6}\" + whitespace +\n\t\"?|\\\\\\\\([^\\\\r\\\\n\\\\f])\", \"g\" ),\n\tfunescape = function( escape, nonHex ) {\n\t\tvar high = \"0x\" + escape.slice( 1 ) - 0x10000;\n\n\t\tif ( nonHex ) {\n\n\t\t\t// Strip the backslash prefix from a non-hex escape sequence\n\t\t\treturn nonHex;\n\t\t}\n\n\t\t// Replace a hexadecimal escape sequence with the encoded Unicode code point\n\t\t// Support: IE <=11+\n\t\t// For values outside the Basic Multilingual Plane (BMP), manually construct a\n\t\t// surrogate pair\n\t\treturn high < 0 ?\n\t\t\tString.fromCharCode( high + 0x10000 ) :\n\t\t\tString.fromCharCode( high >> 10 | 0xD800, high & 0x3FF | 0xDC00 );\n\t};\n\nfunction unescapeSelector( sel ) {\n\treturn sel.replace( runescape, funescape );\n}\n\nfunction selectorError( msg ) {\n\tjQuery.error( \"Syntax error, unrecognized expression: \" + msg );\n}\n\nvar rcomma = new RegExp( \"^\" + whitespace + \"*,\" + whitespace + \"*\" );\n\nvar tokenCache = createCache();\n\nfunction tokenize( selector, parseOnly ) {\n\tvar matched, match, tokens, type,\n\t\tsoFar, groups, preFilters,\n\t\tcached = tokenCache[ selector + \" \" ];\n\n\tif ( cached ) {\n\t\treturn parseOnly ? 0 : cached.slice( 0 );\n\t}\n\n\tsoFar = selector;\n\tgroups = [];\n\tpreFilters = jQuery.expr.preFilter;\n\n\twhile ( soFar ) {\n\n\t\t// Comma and first run\n\t\tif ( !matched || ( match = rcomma.exec( soFar ) ) ) {\n\t\t\tif ( match ) {\n\n\t\t\t\t// Don't consume trailing commas as valid\n\t\t\t\tsoFar = soFar.slice( match[ 0 ].length ) || soFar;\n\t\t\t}\n\t\t\tgroups.push( ( tokens = [] ) );\n\t\t}\n\n\t\tmatched = false;\n\n\t\t// Combinators\n\t\tif ( ( match = rleadingCombinator.exec( soFar ) ) ) {\n\t\t\tmatched = match.shift();\n\t\t\ttokens.push( {\n\t\t\t\tvalue: matched,\n\n\t\t\t\t// Cast descendant combinators to space\n\t\t\t\ttype: match[ 0 ].replace( rtrimCSS, \" \" )\n\t\t\t} );\n\t\t\tsoFar = soFar.slice( matched.length );\n\t\t}\n\n\t\t// Filters\n\t\tfor ( type in filterMatchExpr ) {\n\t\t\tif ( ( match = jQuery.expr.match[ type ].exec( soFar ) ) && ( !preFilters[ type ] ||\n\t\t\t\t( match = preFilters[ type ]( match ) ) ) ) {\n\t\t\t\tmatched = match.shift();\n\t\t\t\ttokens.push( {\n\t\t\t\t\tvalue: matched,\n\t\t\t\t\ttype: type,\n\t\t\t\t\tmatches: match\n\t\t\t\t} );\n\t\t\t\tsoFar = soFar.slice( matched.length );\n\t\t\t}\n\t\t}\n\n\t\tif ( !matched ) {\n\t\t\tbreak;\n\t\t}\n\t}\n\n\t// Return the length of the invalid excess\n\t// if we're just parsing\n\t// Otherwise, throw an error or return tokens\n\tif ( parseOnly ) {\n\t\treturn soFar.length;\n\t}\n\n\treturn soFar ?\n\t\tselectorError( selector ) :\n\n\t\t// Cache the tokens\n\t\ttokenCache( selector, groups ).slice( 0 );\n}\n\nvar preFilter = {\n\tATTR: function( match ) {\n\t\tmatch[ 1 ] = unescapeSelector( match[ 1 ] );\n\n\t\t// Move the given value to match[3] whether quoted or unquoted\n\t\tmatch[ 3 ] = unescapeSelector( match[ 3 ] || match[ 4 ] || match[ 5 ] || \"\" );\n\n\t\tif ( match[ 2 ] === \"~=\" ) {\n\t\t\tmatch[ 3 ] = \" \" + match[ 3 ] + \" \";\n\t\t}\n\n\t\treturn match.slice( 0, 4 );\n\t},\n\n\tCHILD: function( match ) {\n\n\t\t/* matches from filterMatchExpr[\"CHILD\"]\n\t\t\t1 type (only|nth|...)\n\t\t\t2 what (child|of-type)\n\t\t\t3 argument (even|odd|\\d*|\\d*n([+-]\\d+)?|...)\n\t\t\t4 xn-component of xn+y argument ([+-]?\\d*n|)\n\t\t\t5 sign of xn-component\n\t\t\t6 x of xn-component\n\t\t\t7 sign of y-component\n\t\t\t8 y of y-component\n\t\t*/\n\t\tmatch[ 1 ] = match[ 1 ].toLowerCase();\n\n\t\tif ( match[ 1 ].slice( 0, 3 ) === \"nth\" ) {\n\n\t\t\t// nth-* requires argument\n\t\t\tif ( !match[ 3 ] ) {\n\t\t\t\tselectorError( match[ 0 ] );\n\t\t\t}\n\n\t\t\t// numeric x and y parameters for jQuery.expr.filter.CHILD\n\t\t\t// remember that false/true cast respectively to 0/1\n\t\t\tmatch[ 4 ] = +( match[ 4 ] ?\n\t\t\t\tmatch[ 5 ] + ( match[ 6 ] || 1 ) :\n\t\t\t\t2 * ( match[ 3 ] === \"even\" || match[ 3 ] === \"odd\" )\n\t\t\t);\n\t\t\tmatch[ 5 ] = +( ( match[ 7 ] + match[ 8 ] ) || match[ 3 ] === \"odd\" );\n\n\t\t// other types prohibit arguments\n\t\t} else if ( match[ 3 ] ) {\n\t\t\tselectorError( match[ 0 ] );\n\t\t}\n\n\t\treturn match;\n\t},\n\n\tPSEUDO: function( match ) {\n\t\tvar excess,\n\t\t\tunquoted = !match[ 6 ] && match[ 2 ];\n\n\t\tif ( filterMatchExpr.CHILD.test( match[ 0 ] ) ) {\n\t\t\treturn null;\n\t\t}\n\n\t\t// Accept quoted arguments as-is\n\t\tif ( match[ 3 ] ) {\n\t\t\tmatch[ 2 ] = match[ 4 ] || match[ 5 ] || \"\";\n\n\t\t// Strip excess characters from unquoted arguments\n\t\t} else if ( unquoted && rpseudo.test( unquoted ) &&\n\n\t\t\t// Get excess from tokenize (recursively)\n\t\t\t( excess = tokenize( unquoted, true ) ) &&\n\n\t\t\t// advance to the next closing parenthesis\n\t\t\t( excess = unquoted.indexOf( \")\", unquoted.length - excess ) -\n\t\t\t\tunquoted.length ) ) {\n\n\t\t\t// excess is a negative index\n\t\t\tmatch[ 0 ] = match[ 0 ].slice( 0, excess );\n\t\t\tmatch[ 2 ] = unquoted.slice( 0, excess );\n\t\t}\n\n\t\t// Return only captures needed by the pseudo filter method (type and argument)\n\t\treturn match.slice( 0, 3 );\n\t}\n};\n\nfunction toSelector( tokens ) {\n\tvar i = 0,\n\t\tlen = tokens.length,\n\t\tselector = \"\";\n\tfor ( ; i < len; i++ ) {\n\t\tselector += tokens[ i ].value;\n\t}\n\treturn selector;\n}\n\n// Multifunctional method to get and set values of a collection\n// The value/s can optionally be executed if it's a function\nfunction access( elems, fn, key, value, chainable, emptyGet, raw ) {\n\tvar i = 0,\n\t\tlen = elems.length,\n\t\tbulk = key == null;\n\n\t// Sets many values\n\tif ( toType( key ) === \"object\" ) {\n\t\tchainable = true;\n\t\tfor ( i in key ) {\n\t\t\taccess( elems, fn, i, key[ i ], true, emptyGet, raw );\n\t\t}\n\n\t// Sets one value\n\t} else if ( value !== undefined ) {\n\t\tchainable = true;\n\n\t\tif ( typeof value !== \"function\" ) {\n\t\t\traw = true;\n\t\t}\n\n\t\tif ( bulk ) {\n\n\t\t\t// Bulk operations run against the entire set\n\t\t\tif ( raw ) {\n\t\t\t\tfn.call( elems, value );\n\t\t\t\tfn = null;\n\n\t\t\t// ...except when executing function values\n\t\t\t} else {\n\t\t\t\tbulk = fn;\n\t\t\t\tfn = function( elem, _key, value ) {\n\t\t\t\t\treturn bulk.call( jQuery( elem ), value );\n\t\t\t\t};\n\t\t\t}\n\t\t}\n\n\t\tif ( fn ) {\n\t\t\tfor ( ; i < len; i++ ) {\n\t\t\t\tfn(\n\t\t\t\t\telems[ i ], key, raw ?\n\t\t\t\t\t\tvalue :\n\t\t\t\t\t\tvalue.call( elems[ i ], i, fn( elems[ i ], key ) )\n\t\t\t\t);\n\t\t\t}\n\t\t}\n\t}\n\n\tif ( chainable ) {\n\t\treturn elems;\n\t}\n\n\t// Gets\n\tif ( bulk ) {\n\t\treturn fn.call( elems );\n\t}\n\n\treturn len ? fn( elems[ 0 ], key ) : emptyGet;\n}\n\n// Only count HTML whitespace\n// Other whitespace should count in values\n// https://infra.spec.whatwg.org/#ascii-whitespace\nvar rnothtmlwhite = /[^\\x20\\t\\r\\n\\f]+/g;\n\njQuery.fn.extend( {\n\tattr: function( name, value ) {\n\t\treturn access( this, jQuery.attr, name, value, arguments.length > 1 );\n\t},\n\n\tremoveAttr: function( name ) {\n\t\treturn this.each( function() {\n\t\t\tjQuery.removeAttr( this, name );\n\t\t} );\n\t}\n} );\n\njQuery.extend( {\n\tattr: function( elem, name, value ) {\n\t\tvar ret, hooks,\n\t\t\tnType = elem.nodeType;\n\n\t\t// Don't get/set attributes on text, comment and attribute nodes\n\t\tif ( nType === 3 || nType === 8 || nType === 2 ) {\n\t\t\treturn;\n\t\t}\n\n\t\t// Fallback to prop when attributes are not supported\n\t\tif ( typeof elem.getAttribute === \"undefined\" ) {\n\t\t\treturn jQuery.prop( elem, name, value );\n\t\t}\n\n\t\t// Attribute hooks are determined by the lowercase version\n\t\t// Grab necessary hook if one is defined\n\t\tif ( nType !== 1 || !jQuery.isXMLDoc( elem ) ) {\n\t\t\thooks = jQuery.attrHooks[ name.toLowerCase() ];\n\t\t}\n\n\t\tif ( value !== undefined ) {\n\t\t\tif ( value === null ||\n\n\t\t\t\t// For compat with previous handling of boolean attributes,\n\t\t\t\t// remove when `false` passed. For ARIA attributes -\n\t\t\t\t// many of which recognize a `\"false\"` value - continue to\n\t\t\t\t// set the `\"false\"` value as jQuery <4 did.\n\t\t\t\t( value === false && name.toLowerCase().indexOf( \"aria-\" ) !== 0 ) ) {\n\n\t\t\t\tjQuery.removeAttr( elem, name );\n\t\t\t\treturn;\n\t\t\t}\n\n\t\t\tif ( hooks && \"set\" in hooks &&\n\t\t\t\t( ret = hooks.set( elem, value, name ) ) !== undefined ) {\n\t\t\t\treturn ret;\n\t\t\t}\n\n\t\t\telem.setAttribute( name, value );\n\t\t\treturn value;\n\t\t}\n\n\t\tif ( hooks && \"get\" in hooks && ( ret = hooks.get( elem, name ) ) !== null ) {\n\t\t\treturn ret;\n\t\t}\n\n\t\tret = elem.getAttribute( name );\n\n\t\t// Non-existent attributes return null, we normalize to undefined\n\t\treturn ret == null ? undefined : ret;\n\t},\n\n\tattrHooks: {},\n\n\tremoveAttr: function( elem, value ) {\n\t\tvar name,\n\t\t\ti = 0,\n\n\t\t\t// Attribute names can contain non-HTML whitespace characters\n\t\t\t// https://html.spec.whatwg.org/multipage/syntax.html#attributes-2\n\t\t\tattrNames = value && value.match( rnothtmlwhite );\n\n\t\tif ( attrNames && elem.nodeType === 1 ) {\n\t\t\twhile ( ( name = attrNames[ i++ ] ) ) {\n\t\t\t\telem.removeAttribute( name );\n\t\t\t}\n\t\t}\n\t}\n} );\n\n// Support: IE <=11+\n// An input loses its value after becoming a radio\nif ( isIE ) {\n\tjQuery.attrHooks.type = {\n\t\tset: function( elem, value ) {\n\t\t\tif ( value === \"radio\" && nodeName( elem, \"input\" ) ) {\n\t\t\t\tvar val = elem.value;\n\t\t\t\telem.setAttribute( \"type\", value );\n\t\t\t\tif ( val ) {\n\t\t\t\t\telem.value = val;\n\t\t\t\t}\n\t\t\t\treturn value;\n\t\t\t}\n\t\t}\n\t};\n}\n\n// CSS string/identifier serialization\n// https://drafts.csswg.org/cssom/#common-serializing-idioms\nvar rcssescape = /([\\0-\\x1f\\x7f]|^-?\\d)|^-$|[^\\x80-\\uFFFF\\w-]/g;\n\nfunction fcssescape( ch, asCodePoint ) {\n\tif ( asCodePoint ) {\n\n\t\t// U+0000 NULL becomes U+FFFD REPLACEMENT CHARACTER\n\t\tif ( ch === \"\\0\" ) {\n\t\t\treturn \"\\uFFFD\";\n\t\t}\n\n\t\t// Control characters and (dependent upon position) numbers get escaped as code points\n\t\treturn ch.slice( 0, -1 ) + \"\\\\\" + ch.charCodeAt( ch.length - 1 ).toString( 16 ) + \" \";\n\t}\n\n\t// Other potentially-special ASCII characters get backslash-escaped\n\treturn \"\\\\\" + ch;\n}\n\njQuery.escapeSelector = function( sel ) {\n\treturn ( sel + \"\" ).replace( rcssescape, fcssescape );\n};\n\nvar sort = arr.sort;\n\nvar splice = arr.splice;\n\nvar hasDuplicate;\n\n// Document order sorting\nfunction sortOrder( a, b ) {\n\n\t// Flag for duplicate removal\n\tif ( a === b ) {\n\t\thasDuplicate = true;\n\t\treturn 0;\n\t}\n\n\t// Sort on method existence if only one input has compareDocumentPosition\n\tvar compare = !a.compareDocumentPosition - !b.compareDocumentPosition;\n\tif ( compare ) {\n\t\treturn compare;\n\t}\n\n\t// Calculate position if both inputs belong to the same document\n\t// Support: IE 11+\n\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t// two documents; shallow comparisons work.\n\t// eslint-disable-next-line eqeqeq\n\tcompare = ( a.ownerDocument || a ) == ( b.ownerDocument || b ) ?\n\t\ta.compareDocumentPosition( b ) :\n\n\t\t// Otherwise we know they are disconnected\n\t\t1;\n\n\t// Disconnected nodes\n\tif ( compare & 1 ) {\n\n\t\t// Choose the first element that is related to the document\n\t\t// Support: IE 11+\n\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t// two documents; shallow comparisons work.\n\t\t// eslint-disable-next-line eqeqeq\n\t\tif ( a == document$1 || a.ownerDocument == document$1 &&\n\t\t\tjQuery.contains( document$1, a ) ) {\n\t\t\treturn -1;\n\t\t}\n\n\t\t// Support: IE 11+\n\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t// two documents; shallow comparisons work.\n\t\t// eslint-disable-next-line eqeqeq\n\t\tif ( b == document$1 || b.ownerDocument == document$1 &&\n\t\t\tjQuery.contains( document$1, b ) ) {\n\t\t\treturn 1;\n\t\t}\n\n\t\t// Maintain original order\n\t\treturn 0;\n\t}\n\n\treturn compare & 4 ? -1 : 1;\n}\n\n/**\n * Document sorting and removing duplicates\n * @param {ArrayLike} results\n */\njQuery.uniqueSort = function( results ) {\n\tvar elem,\n\t\tduplicates = [],\n\t\tj = 0,\n\t\ti = 0;\n\n\thasDuplicate = false;\n\n\tsort.call( results, sortOrder );\n\n\tif ( hasDuplicate ) {\n\t\twhile ( ( elem = results[ i++ ] ) ) {\n\t\t\tif ( elem === results[ i ] ) {\n\t\t\t\tj = duplicates.push( i );\n\t\t\t}\n\t\t}\n\t\twhile ( j-- ) {\n\t\t\tsplice.call( results, duplicates[ j ], 1 );\n\t\t}\n\t}\n\n\treturn results;\n};\n\njQuery.fn.uniqueSort = function() {\n\treturn this.pushStack( jQuery.uniqueSort( slice.apply( this ) ) );\n};\n\nvar i,\n\toutermostContext,\n\n\t// Local document vars\n\tdocument,\n\tdocumentElement,\n\tdocumentIsHTML,\n\n\t// Instance-specific data\n\tdirruns = 0,\n\tdone = 0,\n\tclassCache = createCache(),\n\tcompilerCache = createCache(),\n\tnonnativeSelectorCache = createCache(),\n\n\t// Regular expressions\n\n\t// Leading and non-escaped trailing whitespace, capturing some non-whitespace characters preceding the latter\n\trwhitespace = new RegExp( whitespace + \"+\", \"g\" ),\n\n\tridentifier = new RegExp( \"^\" + identifier + \"$\" ),\n\n\tmatchExpr = jQuery.extend( {\n\n\t\t// For use in libraries implementing .is()\n\t\t// We use this for POS matching in `select`\n\t\tneedsContext: new RegExp( \"^\" + whitespace +\n\t\t\t\"*[>+~]|:(even|odd|eq|gt|lt|nth|first|last)(?:\\\\(\" + whitespace +\n\t\t\t\"*((?:-\\\\d)?\\\\d*)\" + whitespace + \"*\\\\)|)(?=[^-]|$)\", \"i\" )\n\t}, filterMatchExpr ),\n\n\trinputs = /^(?:input|select|textarea|button)$/i,\n\trheader = /^h\\d$/i,\n\n\t// Easily-parseable/retrievable ID or TAG or CLASS selectors\n\trquickExpr$1 = /^(?:#([\\w-]+)|(\\w+)|\\.([\\w-]+))$/,\n\n\t// Used for iframes; see `setDocument`.\n\t// Support: IE 9 - 11+\n\t// Removing the function wrapper causes a \"Permission Denied\"\n\t// error in IE.\n\tunloadHandler = function() {\n\t\tsetDocument();\n\t},\n\n\tinDisabledFieldset = addCombinator(\n\t\tfunction( elem ) {\n\t\t\treturn elem.disabled === true && nodeName( elem, \"fieldset\" );\n\t\t},\n\t\t{ dir: \"parentNode\", next: \"legend\" }\n\t);\n\nfunction find( selector, context, results, seed ) {\n\tvar m, i, elem, nid, match, groups, newSelector,\n\t\tnewContext = context && context.ownerDocument,\n\n\t\t// nodeType defaults to 9, since context defaults to document\n\t\tnodeType = context ? context.nodeType : 9;\n\n\tresults = results || [];\n\n\t// Return early from calls with invalid selector or context\n\tif ( typeof selector !== \"string\" || !selector ||\n\t\tnodeType !== 1 && nodeType !== 9 && nodeType !== 11 ) {\n\n\t\treturn results;\n\t}\n\n\t// Try to shortcut find operations (as opposed to filters) in HTML documents\n\tif ( !seed ) {\n\t\tsetDocument( context );\n\t\tcontext = context || document;\n\n\t\tif ( documentIsHTML ) {\n\n\t\t\t// If the selector is sufficiently simple, try using a \"get*By*\" DOM method\n\t\t\t// (excepting DocumentFragment context, where the methods don't exist)\n\t\t\tif ( nodeType !== 11 && ( match = rquickExpr$1.exec( selector ) ) ) {\n\n\t\t\t\t// ID selector\n\t\t\t\tif ( ( m = match[ 1 ] ) ) {\n\n\t\t\t\t\t// Document context\n\t\t\t\t\tif ( nodeType === 9 ) {\n\t\t\t\t\t\tif ( ( elem = context.getElementById( m ) ) ) {\n\t\t\t\t\t\t\tpush.call( results, elem );\n\t\t\t\t\t\t}\n\t\t\t\t\t\treturn results;\n\n\t\t\t\t\t// Element context\n\t\t\t\t\t} else {\n\t\t\t\t\t\tif ( newContext && ( elem = newContext.getElementById( m ) ) &&\n\t\t\t\t\t\t\tjQuery.contains( context, elem ) ) {\n\n\t\t\t\t\t\t\tpush.call( results, elem );\n\t\t\t\t\t\t\treturn results;\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t// Type selector\n\t\t\t\t} else if ( match[ 2 ] ) {\n\t\t\t\t\tpush.apply( results, context.getElementsByTagName( selector ) );\n\t\t\t\t\treturn results;\n\n\t\t\t\t// Class selector\n\t\t\t\t} else if ( ( m = match[ 3 ] ) && context.getElementsByClassName ) {\n\t\t\t\t\tpush.apply( results, context.getElementsByClassName( m ) );\n\t\t\t\t\treturn results;\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// Take advantage of querySelectorAll\n\t\t\tif ( !nonnativeSelectorCache[ selector + \" \" ] &&\n\t\t\t\t( !rbuggyQSA || !rbuggyQSA.test( selector ) ) ) {\n\n\t\t\t\tnewSelector = selector;\n\t\t\t\tnewContext = context;\n\n\t\t\t\t// qSA considers elements outside a scoping root when evaluating child or\n\t\t\t\t// descendant combinators, which is not what we want.\n\t\t\t\t// In such cases, we work around the behavior by prefixing every selector in the\n\t\t\t\t// list with an ID selector referencing the scope context.\n\t\t\t\t// The technique has to be used as well when a leading combinator is used\n\t\t\t\t// as such selectors are not recognized by querySelectorAll.\n\t\t\t\t// Thanks to Andrew Dupont for this technique.\n\t\t\t\tif ( nodeType === 1 &&\n\t\t\t\t\t( rdescend.test( selector ) || rleadingCombinator.test( selector ) ) ) {\n\n\t\t\t\t\t// Expand context for sibling selectors\n\t\t\t\t\tnewContext = rsibling.test( selector ) &&\n\t\t\t\t\t\ttestContext( context.parentNode ) ||\n\t\t\t\t\t\tcontext;\n\n\t\t\t\t\t// Outside of IE, if we're not changing the context we can\n\t\t\t\t\t// use :scope instead of an ID.\n\t\t\t\t\t// Support: IE 11+\n\t\t\t\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t\t\t\t// two documents; shallow comparisons work.\n\t\t\t\t\t// eslint-disable-next-line eqeqeq\n\t\t\t\t\tif ( newContext != context || isIE ) {\n\n\t\t\t\t\t\t// Capture the context ID, setting it first if necessary\n\t\t\t\t\t\tif ( ( nid = context.getAttribute( \"id\" ) ) ) {\n\t\t\t\t\t\t\tnid = jQuery.escapeSelector( nid );\n\t\t\t\t\t\t} else {\n\t\t\t\t\t\t\tcontext.setAttribute( \"id\", ( nid = jQuery.expando ) );\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t\t// Prefix every selector in the list\n\t\t\t\t\tgroups = tokenize( selector );\n\t\t\t\t\ti = groups.length;\n\t\t\t\t\twhile ( i-- ) {\n\t\t\t\t\t\tgroups[ i ] = ( nid ? \"#\" + nid : \":scope\" ) + \" \" +\n\t\t\t\t\t\t\ttoSelector( groups[ i ] );\n\t\t\t\t\t}\n\t\t\t\t\tnewSelector = groups.join( \",\" );\n\t\t\t\t}\n\n\t\t\t\ttry {\n\t\t\t\t\tpush.apply( results,\n\t\t\t\t\t\tnewContext.querySelectorAll( newSelector )\n\t\t\t\t\t);\n\t\t\t\t\treturn results;\n\t\t\t\t} catch ( qsaError ) {\n\t\t\t\t\tnonnativeSelectorCache( selector, true );\n\t\t\t\t} finally {\n\t\t\t\t\tif ( nid === jQuery.expando ) {\n\t\t\t\t\t\tcontext.removeAttribute( \"id\" );\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\t// All others\n\treturn select( selector.replace( rtrimCSS, \"$1\" ), context, results, seed );\n}\n\n/**\n * Mark a function for special use by jQuery selector module\n * @param {Function} fn The function to mark\n */\nfunction markFunction( fn ) {\n\tfn[ jQuery.expando ] = true;\n\treturn fn;\n}\n\n/**\n * Returns a function to use in pseudos for input types\n * @param {String} type\n */\nfunction createInputPseudo( type ) {\n\treturn function( elem ) {\n\t\treturn nodeName( elem, \"input\" ) && elem.type === type;\n\t};\n}\n\n/**\n * Returns a function to use in pseudos for buttons\n * @param {String} type\n */\nfunction createButtonPseudo( type ) {\n\treturn function( elem ) {\n\t\treturn ( nodeName( elem, \"input\" ) || nodeName( elem, \"button\" ) ) &&\n\t\t\telem.type === type;\n\t};\n}\n\n/**\n * Returns a function to use in pseudos for :enabled/:disabled\n * @param {Boolean} disabled true for :disabled; false for :enabled\n */\nfunction createDisabledPseudo( disabled ) {\n\n\t// Known :disabled false positives: fieldset[disabled] > legend:nth-of-type(n+2) :can-disable\n\treturn function( elem ) {\n\n\t\t// Only certain elements can match :enabled or :disabled\n\t\t// https://html.spec.whatwg.org/multipage/scripting.html#selector-enabled\n\t\t// https://html.spec.whatwg.org/multipage/scripting.html#selector-disabled\n\t\tif ( \"form\" in elem ) {\n\n\t\t\t// Check for inherited disabledness on relevant non-disabled elements:\n\t\t\t// * listed form-associated elements in a disabled fieldset\n\t\t\t// https://html.spec.whatwg.org/multipage/forms.html#category-listed\n\t\t\t// https://html.spec.whatwg.org/multipage/forms.html#concept-fe-disabled\n\t\t\t// * option elements in a disabled optgroup\n\t\t\t// https://html.spec.whatwg.org/multipage/forms.html#concept-option-disabled\n\t\t\t// All such elements have a \"form\" property.\n\t\t\tif ( elem.parentNode && elem.disabled === false ) {\n\n\t\t\t\t// Option elements defer to a parent optgroup if present\n\t\t\t\tif ( \"label\" in elem ) {\n\t\t\t\t\tif ( \"label\" in elem.parentNode ) {\n\t\t\t\t\t\treturn elem.parentNode.disabled === disabled;\n\t\t\t\t\t} else {\n\t\t\t\t\t\treturn elem.disabled === disabled;\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\t// Support: IE 6 - 11+\n\t\t\t\t// Use the isDisabled shortcut property to check for disabled fieldset ancestors\n\t\t\t\treturn elem.isDisabled === disabled ||\n\n\t\t\t\t\t// Where there is no isDisabled, check manually\n\t\t\t\t\telem.isDisabled !== !disabled &&\n\t\t\t\t\t\tinDisabledFieldset( elem ) === disabled;\n\t\t\t}\n\n\t\t\treturn elem.disabled === disabled;\n\n\t\t// Try to winnow out elements that can't be disabled before trusting the disabled property.\n\t\t// Some victims get caught in our net (label, legend, menu, track), but it shouldn't\n\t\t// even exist on them, let alone have a boolean value.\n\t\t} else if ( \"label\" in elem ) {\n\t\t\treturn elem.disabled === disabled;\n\t\t}\n\n\t\t// Remaining elements are neither :enabled nor :disabled\n\t\treturn false;\n\t};\n}\n\n/**\n * Returns a function to use in pseudos for positionals\n * @param {Function} fn\n */\nfunction createPositionalPseudo( fn ) {\n\treturn markFunction( function( argument ) {\n\t\targument = +argument;\n\t\treturn markFunction( function( seed, matches ) {\n\t\t\tvar j,\n\t\t\t\tmatchIndexes = fn( [], seed.length, argument ),\n\t\t\t\ti = matchIndexes.length;\n\n\t\t\t// Match elements found at the specified indexes\n\t\t\twhile ( i-- ) {\n\t\t\t\tif ( seed[ ( j = matchIndexes[ i ] ) ] ) {\n\t\t\t\t\tseed[ j ] = !( matches[ j ] = seed[ j ] );\n\t\t\t\t}\n\t\t\t}\n\t\t} );\n\t} );\n}\n\n/**\n * Sets document-related variables once based on the current document\n * @param {Element|Object} [node] An element or document object to use to set the document\n */\nfunction setDocument( node ) {\n\tvar subWindow,\n\t\tdoc = node ? node.ownerDocument || node : document$1;\n\n\t// Return early if doc is invalid or already selected\n\t// Support: IE 11+\n\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t// two documents; shallow comparisons work.\n\t// eslint-disable-next-line eqeqeq\n\tif ( doc == document || doc.nodeType !== 9 ) {\n\t\treturn;\n\t}\n\n\t// Update global variables\n\tdocument = doc;\n\tdocumentElement = document.documentElement;\n\tdocumentIsHTML = !jQuery.isXMLDoc( document );\n\n\t// Support: IE 9 - 11+\n\t// Accessing iframe documents after unload throws \"permission denied\" errors (see trac-13936)\n\t// Support: IE 11+\n\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t// two documents; shallow comparisons work.\n\t// eslint-disable-next-line eqeqeq\n\tif ( isIE && document$1 != document &&\n\t\t( subWindow = document.defaultView ) && subWindow.top !== subWindow ) {\n\t\tsubWindow.addEventListener( \"unload\", unloadHandler );\n\t}\n}\n\nfind.matches = function( expr, elements ) {\n\treturn find( expr, null, null, elements );\n};\n\nfind.matchesSelector = function( elem, expr ) {\n\tsetDocument( elem );\n\n\tif ( documentIsHTML &&\n\t\t!nonnativeSelectorCache[ expr + \" \" ] &&\n\t\t( !rbuggyQSA || !rbuggyQSA.test( expr ) ) ) {\n\n\t\ttry {\n\t\t\treturn matches.call( elem, expr );\n\t\t} catch ( e ) {\n\t\t\tnonnativeSelectorCache( expr, true );\n\t\t}\n\t}\n\n\treturn find( expr, document, null, [ elem ] ).length > 0;\n};\n\njQuery.expr = {\n\n\t// Can be adjusted by the user\n\tcacheLength: 50,\n\n\tcreatePseudo: markFunction,\n\n\tmatch: matchExpr,\n\n\tfind: {\n\t\tID: function( id, context ) {\n\t\t\tif ( typeof context.getElementById !== \"undefined\" && documentIsHTML ) {\n\t\t\t\tvar elem = context.getElementById( id );\n\t\t\t\treturn elem ? [ elem ] : [];\n\t\t\t}\n\t\t},\n\n\t\tTAG: function( tag, context ) {\n\t\t\tif ( typeof context.getElementsByTagName !== \"undefined\" ) {\n\t\t\t\treturn context.getElementsByTagName( tag );\n\n\t\t\t\t// DocumentFragment nodes don't have gEBTN\n\t\t\t} else {\n\t\t\t\treturn context.querySelectorAll( tag );\n\t\t\t}\n\t\t},\n\n\t\tCLASS: function( className, context ) {\n\t\t\tif ( typeof context.getElementsByClassName !== \"undefined\" && documentIsHTML ) {\n\t\t\t\treturn context.getElementsByClassName( className );\n\t\t\t}\n\t\t}\n\t},\n\n\trelative: {\n\t\t\">\": { dir: \"parentNode\", first: true },\n\t\t\" \": { dir: \"parentNode\" },\n\t\t\"+\": { dir: \"previousSibling\", first: true },\n\t\t\"~\": { dir: \"previousSibling\" }\n\t},\n\n\tpreFilter: preFilter,\n\n\tfilter: {\n\t\tID: function( id ) {\n\t\t\tvar attrId = unescapeSelector( id );\n\t\t\treturn function( elem ) {\n\t\t\t\treturn elem.getAttribute( \"id\" ) === attrId;\n\t\t\t};\n\t\t},\n\n\t\tTAG: function( nodeNameSelector ) {\n\t\t\tvar expectedNodeName = unescapeSelector( nodeNameSelector ).toLowerCase();\n\t\t\treturn nodeNameSelector === \"*\" ?\n\n\t\t\t\tfunction() {\n\t\t\t\t\treturn true;\n\t\t\t\t} :\n\n\t\t\t\tfunction( elem ) {\n\t\t\t\t\treturn nodeName( elem, expectedNodeName );\n\t\t\t\t};\n\t\t},\n\n\t\tCLASS: function( className ) {\n\t\t\tvar pattern = classCache[ className + \" \" ];\n\n\t\t\treturn pattern ||\n\t\t\t\t( pattern = new RegExp( \"(^|\" + whitespace + \")\" + className +\n\t\t\t\t\t\"(\" + whitespace + \"|$)\" ) ) &&\n\t\t\t\tclassCache( className, function( elem ) {\n\t\t\t\t\treturn pattern.test(\n\t\t\t\t\t\ttypeof elem.className === \"string\" && elem.className ||\n\t\t\t\t\t\t\ttypeof elem.getAttribute !== \"undefined\" &&\n\t\t\t\t\t\t\t\telem.getAttribute( \"class\" ) ||\n\t\t\t\t\t\t\t\"\"\n\t\t\t\t\t);\n\t\t\t\t} );\n\t\t},\n\n\t\tATTR: function( name, operator, check ) {\n\t\t\treturn function( elem ) {\n\t\t\t\tvar result = jQuery.attr( elem, name );\n\n\t\t\t\tif ( result == null ) {\n\t\t\t\t\treturn operator === \"!=\";\n\t\t\t\t}\n\t\t\t\tif ( !operator ) {\n\t\t\t\t\treturn true;\n\t\t\t\t}\n\n\t\t\t\tresult += \"\";\n\n\t\t\t\tif ( operator === \"=\" ) {\n\t\t\t\t\treturn result === check;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"!=\" ) {\n\t\t\t\t\treturn result !== check;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"^=\" ) {\n\t\t\t\t\treturn check && result.indexOf( check ) === 0;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"*=\" ) {\n\t\t\t\t\treturn check && result.indexOf( check ) > -1;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"$=\" ) {\n\t\t\t\t\treturn check && result.slice( -check.length ) === check;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"~=\" ) {\n\t\t\t\t\treturn ( \" \" + result.replace( rwhitespace, \" \" ) + \" \" )\n\t\t\t\t\t\t.indexOf( check ) > -1;\n\t\t\t\t}\n\t\t\t\tif ( operator === \"|=\" ) {\n\t\t\t\t\treturn result === check || result.slice( 0, check.length + 1 ) === check + \"-\";\n\t\t\t\t}\n\n\t\t\t\treturn false;\n\t\t\t};\n\t\t},\n\n\t\tCHILD: function( type, what, _argument, first, last ) {\n\t\t\tvar simple = type.slice( 0, 3 ) !== \"nth\",\n\t\t\t\tforward = type.slice( -4 ) !== \"last\",\n\t\t\t\tofType = what === \"of-type\";\n\n\t\t\treturn first === 1 && last === 0 ?\n\n\t\t\t\t// Shortcut for :nth-*(n)\n\t\t\t\tfunction( elem ) {\n\t\t\t\t\treturn !!elem.parentNode;\n\t\t\t\t} :\n\n\t\t\t\tfunction( elem, _context, xml ) {\n\t\t\t\t\tvar cache, outerCache, node, nodeIndex, start,\n\t\t\t\t\t\tdir = simple !== forward ? \"nextSibling\" : \"previousSibling\",\n\t\t\t\t\t\tparent = elem.parentNode,\n\t\t\t\t\t\tname = ofType && elem.nodeName.toLowerCase(),\n\t\t\t\t\t\tuseCache = !xml && !ofType,\n\t\t\t\t\t\tdiff = false;\n\n\t\t\t\t\tif ( parent ) {\n\n\t\t\t\t\t\t// :(first|last|only)-(child|of-type)\n\t\t\t\t\t\tif ( simple ) {\n\t\t\t\t\t\t\twhile ( dir ) {\n\t\t\t\t\t\t\t\tnode = elem;\n\t\t\t\t\t\t\t\twhile ( ( node = node[ dir ] ) ) {\n\t\t\t\t\t\t\t\t\tif ( ofType ?\n\t\t\t\t\t\t\t\t\t\tnodeName( node, name ) :\n\t\t\t\t\t\t\t\t\t\tnode.nodeType === 1 ) {\n\n\t\t\t\t\t\t\t\t\t\treturn false;\n\t\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\t\t// Reverse direction for :only-* (if we haven't yet done so)\n\t\t\t\t\t\t\t\tstart = dir = type === \"only\" && !start && \"nextSibling\";\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\treturn true;\n\t\t\t\t\t\t}\n\n\t\t\t\t\t\tstart = [ forward ? parent.firstChild : parent.lastChild ];\n\n\t\t\t\t\t\t// non-xml :nth-child(...) stores cache data on `parent`\n\t\t\t\t\t\tif ( forward && useCache ) {\n\n\t\t\t\t\t\t\t// Seek `elem` from a previously-cached index\n\t\t\t\t\t\t\touterCache = parent[ jQuery.expando ] ||\n\t\t\t\t\t\t\t\t( parent[ jQuery.expando ] = {} );\n\t\t\t\t\t\t\tcache = outerCache[ type ] || [];\n\t\t\t\t\t\t\tnodeIndex = cache[ 0 ] === dirruns && cache[ 1 ];\n\t\t\t\t\t\t\tdiff = nodeIndex && cache[ 2 ];\n\t\t\t\t\t\t\tnode = nodeIndex && parent.childNodes[ nodeIndex ];\n\n\t\t\t\t\t\t\twhile ( ( node = ++nodeIndex && node && node[ dir ] ||\n\n\t\t\t\t\t\t\t\t// Fallback to seeking `elem` from the start\n\t\t\t\t\t\t\t\t( diff = nodeIndex = 0 ) || start.pop() ) ) {\n\n\t\t\t\t\t\t\t\t// When found, cache indexes on `parent` and break\n\t\t\t\t\t\t\t\tif ( node.nodeType === 1 && ++diff && node === elem ) {\n\t\t\t\t\t\t\t\t\touterCache[ type ] = [ dirruns, nodeIndex, diff ];\n\t\t\t\t\t\t\t\t\tbreak;\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t} else {\n\n\t\t\t\t\t\t\t// Use previously-cached element index if available\n\t\t\t\t\t\t\tif ( useCache ) {\n\t\t\t\t\t\t\t\touterCache = elem[ jQuery.expando ] ||\n\t\t\t\t\t\t\t\t\t( elem[ jQuery.expando ] = {} );\n\t\t\t\t\t\t\t\tcache = outerCache[ type ] || [];\n\t\t\t\t\t\t\t\tnodeIndex = cache[ 0 ] === dirruns && cache[ 1 ];\n\t\t\t\t\t\t\t\tdiff = nodeIndex;\n\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\t// xml :nth-child(...)\n\t\t\t\t\t\t\t// or :nth-last-child(...) or :nth(-last)?-of-type(...)\n\t\t\t\t\t\t\tif ( diff === false ) {\n\n\t\t\t\t\t\t\t\t// Use the same loop as above to seek `elem` from the start\n\t\t\t\t\t\t\t\twhile ( ( node = ++nodeIndex && node && node[ dir ] ||\n\t\t\t\t\t\t\t\t\t( diff = nodeIndex = 0 ) || start.pop() ) ) {\n\n\t\t\t\t\t\t\t\t\tif ( ( ofType ?\n\t\t\t\t\t\t\t\t\t\tnodeName( node, name ) :\n\t\t\t\t\t\t\t\t\t\tnode.nodeType === 1 ) &&\n\t\t\t\t\t\t\t\t\t\t++diff ) {\n\n\t\t\t\t\t\t\t\t\t\t// Cache the index of each encountered element\n\t\t\t\t\t\t\t\t\t\tif ( useCache ) {\n\t\t\t\t\t\t\t\t\t\t\touterCache = node[ jQuery.expando ] ||\n\t\t\t\t\t\t\t\t\t\t\t\t( node[ jQuery.expando ] = {} );\n\t\t\t\t\t\t\t\t\t\t\touterCache[ type ] = [ dirruns, diff ];\n\t\t\t\t\t\t\t\t\t\t}\n\n\t\t\t\t\t\t\t\t\t\tif ( node === elem ) {\n\t\t\t\t\t\t\t\t\t\t\tbreak;\n\t\t\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t\t}\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t}\n\n\t\t\t\t\t\t// Incorporate the offset, then check against cycle size\n\t\t\t\t\t\tdiff -= last;\n\t\t\t\t\t\treturn diff === first || ( diff % first === 0 && diff / first >= 0 );\n\t\t\t\t\t}\n\t\t\t\t};\n\t\t},\n\n\t\tPSEUDO: function( pseudo, argument ) {\n\n\t\t\t// pseudo-class names are case-insensitive\n\t\t\t// https://www.w3.org/TR/selectors/#pseudo-classes\n\t\t\t// Prioritize by case sensitivity in case custom pseudos are added with uppercase letters\n\t\t\t// Remember that setFilters inherits from pseudos\n\t\t\tvar fn = jQuery.expr.pseudos[ pseudo ] ||\n\t\t\t\tjQuery.expr.setFilters[ pseudo.toLowerCase() ] ||\n\t\t\t\tselectorError( \"unsupported pseudo: \" + pseudo );\n\n\t\t\t// The user may use createPseudo to indicate that\n\t\t\t// arguments are needed to create the filter function\n\t\t\t// just as jQuery does\n\t\t\tif ( fn[ jQuery.expando ] ) {\n\t\t\t\treturn fn( argument );\n\t\t\t}\n\n\t\t\treturn fn;\n\t\t}\n\t},\n\n\tpseudos: {\n\n\t\t// Potentially complex pseudos\n\t\tnot: markFunction( function( selector ) {\n\n\t\t\t// Trim the selector passed to compile\n\t\t\t// to avoid treating leading and trailing\n\t\t\t// spaces as combinators\n\t\t\tvar input = [],\n\t\t\t\tresults = [],\n\t\t\t\tmatcher = compile( selector.replace( rtrimCSS, \"$1\" ) );\n\n\t\t\treturn matcher[ jQuery.expando ] ?\n\t\t\t\tmarkFunction( function( seed, matches, _context, xml ) {\n\t\t\t\t\tvar elem,\n\t\t\t\t\t\tunmatched = matcher( seed, null, xml, [] ),\n\t\t\t\t\t\ti = seed.length;\n\n\t\t\t\t\t// Match elements unmatched by `matcher`\n\t\t\t\t\twhile ( i-- ) {\n\t\t\t\t\t\tif ( ( elem = unmatched[ i ] ) ) {\n\t\t\t\t\t\t\tseed[ i ] = !( matches[ i ] = elem );\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t} ) :\n\t\t\t\tfunction( elem, _context, xml ) {\n\t\t\t\t\tinput[ 0 ] = elem;\n\t\t\t\t\tmatcher( input, null, xml, results );\n\n\t\t\t\t\t// Don't keep the element\n\t\t\t\t\t// (see https://github.com/jquery/sizzle/issues/299)\n\t\t\t\t\tinput[ 0 ] = null;\n\t\t\t\t\treturn !results.pop();\n\t\t\t\t};\n\t\t} ),\n\n\t\thas: markFunction( function( selector ) {\n\t\t\treturn function( elem ) {\n\t\t\t\treturn find( selector, elem ).length > 0;\n\t\t\t};\n\t\t} ),\n\n\t\tcontains: markFunction( function( text ) {\n\t\t\ttext = unescapeSelector( text );\n\t\t\treturn function( elem ) {\n\t\t\t\treturn ( elem.textContent || jQuery.text( elem ) ).indexOf( text ) > -1;\n\t\t\t};\n\t\t} ),\n\n\t\t// \"Whether an element is represented by a :lang() selector\n\t\t// is based solely on the element's language value\n\t\t// being equal to the identifier C,\n\t\t// or beginning with the identifier C immediately followed by \"-\".\n\t\t// The matching of C against the element's language value is performed case-insensitively.\n\t\t// The identifier C does not have to be a valid language name.\"\n\t\t// https://www.w3.org/TR/selectors/#lang-pseudo\n\t\tlang: markFunction( function( lang ) {\n\n\t\t\t// lang value must be a valid identifier\n\t\t\tif ( !ridentifier.test( lang || \"\" ) ) {\n\t\t\t\tselectorError( \"unsupported lang: \" + lang );\n\t\t\t}\n\t\t\tlang = unescapeSelector( lang ).toLowerCase();\n\t\t\treturn function( elem ) {\n\t\t\t\tvar elemLang;\n\t\t\t\tdo {\n\t\t\t\t\tif ( ( elemLang = documentIsHTML ?\n\t\t\t\t\t\telem.lang :\n\t\t\t\t\t\telem.getAttribute( \"xml:lang\" ) || elem.getAttribute( \"lang\" ) ) ) {\n\n\t\t\t\t\t\telemLang = elemLang.toLowerCase();\n\t\t\t\t\t\treturn elemLang === lang || elemLang.indexOf( lang + \"-\" ) === 0;\n\t\t\t\t\t}\n\t\t\t\t} while ( ( elem = elem.parentNode ) && elem.nodeType === 1 );\n\t\t\t\treturn false;\n\t\t\t};\n\t\t} ),\n\n\t\t// Miscellaneous\n\t\ttarget: function( elem ) {\n\t\t\tvar hash = window.location && window.location.hash;\n\t\t\treturn hash && hash.slice( 1 ) === elem.id;\n\t\t},\n\n\t\troot: function( elem ) {\n\t\t\treturn elem === documentElement;\n\t\t},\n\n\t\tfocus: function( elem ) {\n\t\t\treturn elem === document.activeElement &&\n\t\t\t\tdocument.hasFocus() &&\n\t\t\t\t!!( elem.type || elem.href || ~elem.tabIndex );\n\t\t},\n\n\t\t// Boolean properties\n\t\tenabled: createDisabledPseudo( false ),\n\t\tdisabled: createDisabledPseudo( true ),\n\n\t\tchecked: function( elem ) {\n\n\t\t\t// In CSS3, :checked should return both checked and selected elements\n\t\t\t// https://www.w3.org/TR/2011/REC-css3-selectors-20110929/#checked\n\t\t\treturn ( nodeName( elem, \"input\" ) && !!elem.checked ) ||\n\t\t\t\t( nodeName( elem, \"option\" ) && !!elem.selected );\n\t\t},\n\n\t\tselected: function( elem ) {\n\n\t\t\t// Support: IE <=11+\n\t\t\t// Accessing the selectedIndex property\n\t\t\t// forces the browser to treat the default option as\n\t\t\t// selected when in an optgroup.\n\t\t\tif ( isIE && elem.parentNode ) {\n\t\t\t\t// eslint-disable-next-line no-unused-expressions\n\t\t\t\telem.parentNode.selectedIndex;\n\t\t\t}\n\n\t\t\treturn elem.selected === true;\n\t\t},\n\n\t\t// Contents\n\t\tempty: function( elem ) {\n\n\t\t\t// https://www.w3.org/TR/selectors/#empty-pseudo\n\t\t\t// :empty is negated by element (1) or content nodes (text: 3; cdata: 4; entity ref: 5),\n\t\t\t// but not by others (comment: 8; processing instruction: 7; etc.)\n\t\t\t// nodeType < 6 works because attributes (2) do not appear as children\n\t\t\tfor ( elem = elem.firstChild; elem; elem = elem.nextSibling ) {\n\t\t\t\tif ( elem.nodeType < 6 ) {\n\t\t\t\t\treturn false;\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn true;\n\t\t},\n\n\t\tparent: function( elem ) {\n\t\t\treturn !jQuery.expr.pseudos.empty( elem );\n\t\t},\n\n\t\t// Element/input types\n\t\theader: function( elem ) {\n\t\t\treturn rheader.test( elem.nodeName );\n\t\t},\n\n\t\tinput: function( elem ) {\n\t\t\treturn rinputs.test( elem.nodeName );\n\t\t},\n\n\t\tbutton: function( elem ) {\n\t\t\treturn nodeName( elem, \"input\" ) && elem.type === \"button\" ||\n\t\t\t\tnodeName( elem, \"button\" );\n\t\t},\n\n\t\ttext: function( elem ) {\n\t\t\treturn nodeName( elem, \"input\" ) && elem.type === \"text\";\n\t\t},\n\n\t\t// Position-in-collection\n\t\tfirst: createPositionalPseudo( function() {\n\t\t\treturn [ 0 ];\n\t\t} ),\n\n\t\tlast: createPositionalPseudo( function( _matchIndexes, length ) {\n\t\t\treturn [ length - 1 ];\n\t\t} ),\n\n\t\teq: createPositionalPseudo( function( _matchIndexes, length, argument ) {\n\t\t\treturn [ argument < 0 ? argument + length : argument ];\n\t\t} ),\n\n\t\teven: createPositionalPseudo( function( matchIndexes, length ) {\n\t\t\tvar i = 0;\n\t\t\tfor ( ; i < length; i += 2 ) {\n\t\t\t\tmatchIndexes.push( i );\n\t\t\t}\n\t\t\treturn matchIndexes;\n\t\t} ),\n\n\t\todd: createPositionalPseudo( function( matchIndexes, length ) {\n\t\t\tvar i = 1;\n\t\t\tfor ( ; i < length; i += 2 ) {\n\t\t\t\tmatchIndexes.push( i );\n\t\t\t}\n\t\t\treturn matchIndexes;\n\t\t} ),\n\n\t\tlt: createPositionalPseudo( function( matchIndexes, length, argument ) {\n\t\t\tvar i;\n\n\t\t\tif ( argument < 0 ) {\n\t\t\t\ti = argument + length;\n\t\t\t} else if ( argument > length ) {\n\t\t\t\ti = length;\n\t\t\t} else {\n\t\t\t\ti = argument;\n\t\t\t}\n\n\t\t\tfor ( ; --i >= 0; ) {\n\t\t\t\tmatchIndexes.push( i );\n\t\t\t}\n\t\t\treturn matchIndexes;\n\t\t} ),\n\n\t\tgt: createPositionalPseudo( function( matchIndexes, length, argument ) {\n\t\t\tvar i = argument < 0 ? argument + length : argument;\n\t\t\tfor ( ; ++i < length; ) {\n\t\t\t\tmatchIndexes.push( i );\n\t\t\t}\n\t\t\treturn matchIndexes;\n\t\t} )\n\t}\n};\n\njQuery.expr.pseudos.nth = jQuery.expr.pseudos.eq;\n\n// Add button/input type pseudos\nfor ( i in { radio: true, checkbox: true, file: true, password: true, image: true } ) {\n\tjQuery.expr.pseudos[ i ] = createInputPseudo( i );\n}\nfor ( i in { submit: true, reset: true } ) {\n\tjQuery.expr.pseudos[ i ] = createButtonPseudo( i );\n}\n\n// Easy API for creating new setFilters\nfunction setFilters() {}\nsetFilters.prototype = jQuery.expr.pseudos;\njQuery.expr.setFilters = new setFilters();\n\nfunction addCombinator( matcher, combinator, base ) {\n\tvar dir = combinator.dir,\n\t\tskip = combinator.next,\n\t\tkey = skip || dir,\n\t\tcheckNonElements = base && key === \"parentNode\",\n\t\tdoneName = done++;\n\n\treturn combinator.first ?\n\n\t\t// Check against closest ancestor/preceding element\n\t\tfunction( elem, context, xml ) {\n\t\t\twhile ( ( elem = elem[ dir ] ) ) {\n\t\t\t\tif ( elem.nodeType === 1 || checkNonElements ) {\n\t\t\t\t\treturn matcher( elem, context, xml );\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn false;\n\t\t} :\n\n\t\t// Check against all ancestor/preceding elements\n\t\tfunction( elem, context, xml ) {\n\t\t\tvar oldCache, outerCache,\n\t\t\t\tnewCache = [ dirruns, doneName ];\n\n\t\t\t// We can't set arbitrary data on XML nodes, so they don't benefit from combinator caching\n\t\t\tif ( xml ) {\n\t\t\t\twhile ( ( elem = elem[ dir ] ) ) {\n\t\t\t\t\tif ( elem.nodeType === 1 || checkNonElements ) {\n\t\t\t\t\t\tif ( matcher( elem, context, xml ) ) {\n\t\t\t\t\t\t\treturn true;\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t} else {\n\t\t\t\twhile ( ( elem = elem[ dir ] ) ) {\n\t\t\t\t\tif ( elem.nodeType === 1 || checkNonElements ) {\n\t\t\t\t\t\touterCache = elem[ jQuery.expando ] || ( elem[ jQuery.expando ] = {} );\n\n\t\t\t\t\t\tif ( skip && nodeName( elem, skip ) ) {\n\t\t\t\t\t\t\telem = elem[ dir ] || elem;\n\t\t\t\t\t\t} else if ( ( oldCache = outerCache[ key ] ) &&\n\t\t\t\t\t\t\toldCache[ 0 ] === dirruns && oldCache[ 1 ] === doneName ) {\n\n\t\t\t\t\t\t\t// Assign to newCache so results back-propagate to previous elements\n\t\t\t\t\t\t\treturn ( newCache[ 2 ] = oldCache[ 2 ] );\n\t\t\t\t\t\t} else {\n\n\t\t\t\t\t\t\t// Reuse newcache so results back-propagate to previous elements\n\t\t\t\t\t\t\touterCache[ key ] = newCache;\n\n\t\t\t\t\t\t\t// A match means we're done; a fail means we have to keep checking\n\t\t\t\t\t\t\tif ( ( newCache[ 2 ] = matcher( elem, context, xml ) ) ) {\n\t\t\t\t\t\t\t\treturn true;\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn false;\n\t\t};\n}\n\nfunction elementMatcher( matchers ) {\n\treturn matchers.length > 1 ?\n\t\tfunction( elem, context, xml ) {\n\t\t\tvar i = matchers.length;\n\t\t\twhile ( i-- ) {\n\t\t\t\tif ( !matchers[ i ]( elem, context, xml ) ) {\n\t\t\t\t\treturn false;\n\t\t\t\t}\n\t\t\t}\n\t\t\treturn true;\n\t\t} :\n\t\tmatchers[ 0 ];\n}\n\nfunction multipleContexts( selector, contexts, results ) {\n\tvar i = 0,\n\t\tlen = contexts.length;\n\tfor ( ; i < len; i++ ) {\n\t\tfind( selector, contexts[ i ], results );\n\t}\n\treturn results;\n}\n\nfunction condense( unmatched, map, filter, context, xml ) {\n\tvar elem,\n\t\tnewUnmatched = [],\n\t\ti = 0,\n\t\tlen = unmatched.length,\n\t\tmapped = map != null;\n\n\tfor ( ; i < len; i++ ) {\n\t\tif ( ( elem = unmatched[ i ] ) ) {\n\t\t\tif ( !filter || filter( elem, context, xml ) ) {\n\t\t\t\tnewUnmatched.push( elem );\n\t\t\t\tif ( mapped ) {\n\t\t\t\t\tmap.push( i );\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\treturn newUnmatched;\n}\n\nfunction setMatcher( preFilter, selector, matcher, postFilter, postFinder, postSelector ) {\n\tif ( postFilter && !postFilter[ jQuery.expando ] ) {\n\t\tpostFilter = setMatcher( postFilter );\n\t}\n\tif ( postFinder && !postFinder[ jQuery.expando ] ) {\n\t\tpostFinder = setMatcher( postFinder, postSelector );\n\t}\n\treturn markFunction( function( seed, results, context, xml ) {\n\t\tvar temp, i, elem, matcherOut,\n\t\t\tpreMap = [],\n\t\t\tpostMap = [],\n\t\t\tpreexisting = results.length,\n\n\t\t\t// Get initial elements from seed or context\n\t\t\telems = seed ||\n\t\t\t\tmultipleContexts( selector || \"*\",\n\t\t\t\t\tcontext.nodeType ? [ context ] : context, [] ),\n\n\t\t\t// Prefilter to get matcher input, preserving a map for seed-results synchronization\n\t\t\tmatcherIn = preFilter && ( seed || !selector ) ?\n\t\t\t\tcondense( elems, preMap, preFilter, context, xml ) :\n\t\t\t\telems;\n\n\t\tif ( matcher ) {\n\n\t\t\t// If we have a postFinder, or filtered seed, or non-seed postFilter\n\t\t\t// or preexisting results,\n\t\t\tmatcherOut = postFinder || ( seed ? preFilter : preexisting || postFilter ) ?\n\n\t\t\t\t// ...intermediate processing is necessary\n\t\t\t\t[] :\n\n\t\t\t\t// ...otherwise use results directly\n\t\t\t\tresults;\n\n\t\t\t// Find primary matches\n\t\t\tmatcher( matcherIn, matcherOut, context, xml );\n\t\t} else {\n\t\t\tmatcherOut = matcherIn;\n\t\t}\n\n\t\t// Apply postFilter\n\t\tif ( postFilter ) {\n\t\t\ttemp = condense( matcherOut, postMap );\n\t\t\tpostFilter( temp, [], context, xml );\n\n\t\t\t// Un-match failing elements by moving them back to matcherIn\n\t\t\ti = temp.length;\n\t\t\twhile ( i-- ) {\n\t\t\t\tif ( ( elem = temp[ i ] ) ) {\n\t\t\t\t\tmatcherOut[ postMap[ i ] ] = !( matcherIn[ postMap[ i ] ] = elem );\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\tif ( seed ) {\n\t\t\tif ( postFinder || preFilter ) {\n\t\t\t\tif ( postFinder ) {\n\n\t\t\t\t\t// Get the final matcherOut by condensing this intermediate into postFinder contexts\n\t\t\t\t\ttemp = [];\n\t\t\t\t\ti = matcherOut.length;\n\t\t\t\t\twhile ( i-- ) {\n\t\t\t\t\t\tif ( ( elem = matcherOut[ i ] ) ) {\n\n\t\t\t\t\t\t\t// Restore matcherIn since elem is not yet a final match\n\t\t\t\t\t\t\ttemp.push( ( matcherIn[ i ] = elem ) );\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tpostFinder( null, ( matcherOut = [] ), temp, xml );\n\t\t\t\t}\n\n\t\t\t\t// Move matched elements from seed to results to keep them synchronized\n\t\t\t\ti = matcherOut.length;\n\t\t\t\twhile ( i-- ) {\n\t\t\t\t\tif ( ( elem = matcherOut[ i ] ) &&\n\t\t\t\t\t\t( temp = postFinder ? indexOf.call( seed, elem ) : preMap[ i ] ) > -1 ) {\n\n\t\t\t\t\t\tseed[ temp ] = !( results[ temp ] = elem );\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t// Add elements to results, through postFinder if defined\n\t\t} else {\n\t\t\tmatcherOut = condense(\n\t\t\t\tmatcherOut === results ?\n\t\t\t\t\tmatcherOut.splice( preexisting, matcherOut.length ) :\n\t\t\t\t\tmatcherOut\n\t\t\t);\n\t\t\tif ( postFinder ) {\n\t\t\t\tpostFinder( null, results, matcherOut, xml );\n\t\t\t} else {\n\t\t\t\tpush.apply( results, matcherOut );\n\t\t\t}\n\t\t}\n\t} );\n}\n\nfunction matcherFromTokens( tokens ) {\n\tvar checkContext, matcher, j,\n\t\tlen = tokens.length,\n\t\tleadingRelative = jQuery.expr.relative[ tokens[ 0 ].type ],\n\t\timplicitRelative = leadingRelative || jQuery.expr.relative[ \" \" ],\n\t\ti = leadingRelative ? 1 : 0,\n\n\t\t// The foundational matcher ensures that elements are reachable from top-level context(s)\n\t\tmatchContext = addCombinator( function( elem ) {\n\t\t\treturn elem === checkContext;\n\t\t}, implicitRelative, true ),\n\t\tmatchAnyContext = addCombinator( function( elem ) {\n\t\t\treturn indexOf.call( checkContext, elem ) > -1;\n\t\t}, implicitRelative, true ),\n\t\tmatchers = [ function( elem, context, xml ) {\n\n\t\t\t// Support: IE 11+\n\t\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t\t// two documents; shallow comparisons work.\n\t\t\t// eslint-disable-next-line eqeqeq\n\t\t\tvar ret = ( !leadingRelative && ( xml || context != outermostContext ) ) || (\n\t\t\t\t( checkContext = context ).nodeType ?\n\t\t\t\t\tmatchContext( elem, context, xml ) :\n\t\t\t\t\tmatchAnyContext( elem, context, xml ) );\n\n\t\t\t// Avoid hanging onto element\n\t\t\t// (see https://github.com/jquery/sizzle/issues/299)\n\t\t\tcheckContext = null;\n\t\t\treturn ret;\n\t\t} ];\n\n\tfor ( ; i < len; i++ ) {\n\t\tif ( ( matcher = jQuery.expr.relative[ tokens[ i ].type ] ) ) {\n\t\t\tmatchers = [ addCombinator( elementMatcher( matchers ), matcher ) ];\n\t\t} else {\n\t\t\tmatcher = jQuery.expr.filter[ tokens[ i ].type ].apply( null, tokens[ i ].matches );\n\n\t\t\t// Return special upon seeing a positional matcher\n\t\t\tif ( matcher[ jQuery.expando ] ) {\n\n\t\t\t\t// Find the next relative operator (if any) for proper handling\n\t\t\t\tj = ++i;\n\t\t\t\tfor ( ; j < len; j++ ) {\n\t\t\t\t\tif ( jQuery.expr.relative[ tokens[ j ].type ] ) {\n\t\t\t\t\t\tbreak;\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t\treturn setMatcher(\n\t\t\t\t\ti > 1 && elementMatcher( matchers ),\n\t\t\t\t\ti > 1 && toSelector(\n\n\t\t\t\t\t\t// If the preceding token was a descendant combinator, insert an implicit any-element `*`\n\t\t\t\t\t\ttokens.slice( 0, i - 1 )\n\t\t\t\t\t\t\t.concat( { value: tokens[ i - 2 ].type === \" \" ? \"*\" : \"\" } )\n\t\t\t\t\t).replace( rtrimCSS, \"$1\" ),\n\t\t\t\t\tmatcher,\n\t\t\t\t\ti < j && matcherFromTokens( tokens.slice( i, j ) ),\n\t\t\t\t\tj < len && matcherFromTokens( ( tokens = tokens.slice( j ) ) ),\n\t\t\t\t\tj < len && toSelector( tokens )\n\t\t\t\t);\n\t\t\t}\n\t\t\tmatchers.push( matcher );\n\t\t}\n\t}\n\n\treturn elementMatcher( matchers );\n}\n\nfunction matcherFromGroupMatchers( elementMatchers, setMatchers ) {\n\tvar bySet = setMatchers.length > 0,\n\t\tbyElement = elementMatchers.length > 0,\n\t\tsuperMatcher = function( seed, context, xml, results, outermost ) {\n\t\t\tvar elem, j, matcher,\n\t\t\t\tmatchedCount = 0,\n\t\t\t\ti = \"0\",\n\t\t\t\tunmatched = seed && [],\n\t\t\t\tsetMatched = [],\n\t\t\t\tcontextBackup = outermostContext,\n\n\t\t\t\t// We must always have either seed elements or outermost context\n\t\t\t\telems = seed || byElement && jQuery.expr.find.TAG( \"*\", outermost ),\n\n\t\t\t\t// Use integer dirruns iff this is the outermost matcher\n\t\t\t\tdirrunsUnique = ( dirruns += contextBackup == null ? 1 : Math.random() || 0.1 );\n\n\t\t\tif ( outermost ) {\n\n\t\t\t\t// Support: IE 11+\n\t\t\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t\t\t// two documents; shallow comparisons work.\n\t\t\t\t// eslint-disable-next-line eqeqeq\n\t\t\t\toutermostContext = context == document || context || outermost;\n\t\t\t}\n\n\t\t\t// Add elements passing elementMatchers directly to results\n\t\t\tfor ( ; ( elem = elems[ i ] ) != null; i++ ) {\n\t\t\t\tif ( byElement && elem ) {\n\t\t\t\t\tj = 0;\n\n\t\t\t\t\t// Support: IE 11+\n\t\t\t\t\t// IE sometimes throws a \"Permission denied\" error when strict-comparing\n\t\t\t\t\t// two documents; shallow comparisons work.\n\t\t\t\t\t// eslint-disable-next-line eqeqeq\n\t\t\t\t\tif ( !context && elem.ownerDocument != document ) {\n\t\t\t\t\t\tsetDocument( elem );\n\t\t\t\t\t\txml = !documentIsHTML;\n\t\t\t\t\t}\n\t\t\t\t\twhile ( ( matcher = elementMatchers[ j++ ] ) ) {\n\t\t\t\t\t\tif ( matcher( elem, context || document, xml ) ) {\n\t\t\t\t\t\t\tpush.call( results, elem );\n\t\t\t\t\t\t\tbreak;\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\t\t\t\t\tif ( outermost ) {\n\t\t\t\t\t\tdirruns = dirrunsUnique;\n\t\t\t\t\t}\n\t\t\t\t}\n\n\t\t\t\t// Track unmatched elements for set filters\n\t\t\t\tif ( bySet ) {\n\n\t\t\t\t\t// They will have gone through all possible matchers\n\t\t\t\t\tif ( ( elem = !matcher && elem ) ) {\n\t\t\t\t\t\tmatchedCount--;\n\t\t\t\t\t}\n\n\t\t\t\t\t// Lengthen the array for every element, matched or not\n\t\t\t\t\tif ( seed ) {\n\t\t\t\t\t\tunmatched.push( elem );\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// `i` is now the count of elements visited above, and adding it to `matchedCount`\n\t\t\t// makes the latter nonnegative.\n\t\t\tmatchedCount += i;\n\n\t\t\t// Apply set filters to unmatched elements\n\t\t\t// NOTE: This can be skipped if there are no unmatched elements (i.e., `matchedCount`\n\t\t\t// equals `i`), unless we didn't visit _any_ elements in the above loop because we have\n\t\t\t// no element matchers and no seed.\n\t\t\t// Incrementing an initially-string \"0\" `i` allows `i` to remain a string only in that\n\t\t\t// case, which will result in a \"00\" `matchedCount` that differs from `i` but is also\n\t\t\t// numerically zero.\n\t\t\tif ( bySet && i !== matchedCount ) {\n\t\t\t\tj = 0;\n\t\t\t\twhile ( ( matcher = setMatchers[ j++ ] ) ) {\n\t\t\t\t\tmatcher( unmatched, setMatched, context, xml );\n\t\t\t\t}\n\n\t\t\t\tif ( seed ) {\n\n\t\t\t\t\t// Reintegrate element matches to eliminate the need for sorting\n\t\t\t\t\tif ( matchedCount > 0 ) {\n\t\t\t\t\t\twhile ( i-- ) {\n\t\t\t\t\t\t\tif ( !( unmatched[ i ] || setMatched[ i ] ) ) {\n\t\t\t\t\t\t\t\tsetMatched[ i ] = pop.call( results );\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t\t// Discard index placeholder values to get only actual matches\n\t\t\t\t\tsetMatched = condense( setMatched );\n\t\t\t\t}\n\n\t\t\t\t// Add matches to results\n\t\t\t\tpush.apply( results, setMatched );\n\n\t\t\t\t// Seedless set matches succeeding multiple successful matchers stipulate sorting\n\t\t\t\tif ( outermost && !seed && setMatched.length > 0 &&\n\t\t\t\t\t( matchedCount + setMatchers.length ) > 1 ) {\n\n\t\t\t\t\tjQuery.uniqueSort( results );\n\t\t\t\t}\n\t\t\t}\n\n\t\t\t// Override manipulation of globals by nested matchers\n\t\t\tif ( outermost ) {\n\t\t\t\tdirruns = dirrunsUnique;\n\t\t\t\toutermostContext = contextBackup;\n\t\t\t}\n\n\t\t\treturn unmatched;\n\t\t};\n\n\treturn bySet ?\n\t\tmarkFunction( superMatcher ) :\n\t\tsuperMatcher;\n}\n\nfunction compile( selector, match /* Internal Use Only */ ) {\n\tvar i,\n\t\tsetMatchers = [],\n\t\telementMatchers = [],\n\t\tcached = compilerCache[ selector + \" \" ];\n\n\tif ( !cached ) {\n\n\t\t// Generate a function of recursive functions that can be used to check each element\n\t\tif ( !match ) {\n\t\t\tmatch = tokenize( selector );\n\t\t}\n\t\ti = match.length;\n\t\twhile ( i-- ) {\n\t\t\tcached = matcherFromTokens( match[ i ] );\n\t\t\tif ( cached[ jQuery.expando ] ) {\n\t\t\t\tsetMatchers.push( cached );\n\t\t\t} else {\n\t\t\t\telementMatchers.push( cached );\n\t\t\t}\n\t\t}\n\n\t\t// Cache the compiled function\n\t\tcached = compilerCache( selector,\n\t\t\tmatcherFromGroupMatchers( elementMatchers, setMatchers ) );\n\n\t\t// Save selector and tokenization\n\t\tcached.selector = selector;\n\t}\n\treturn cached;\n}\n\n/**\n * A low-level selection function that works with jQuery's compiled\n * selector functions\n * @param {String|Function} selector A selector or a pre-compiled\n * selector function built with jQuery selector compile\n * @param {Element} context\n * @param {Array} [results]\n * @param {Array} [seed] A set of elements to match against\n */\nfunction select( selector, context, results, seed ) {\n\tvar i, tokens, token, type, find,\n\t\tcompiled = typeof selector === \"function\" && selector,\n\t\tmatch = !seed && tokenize( ( selector = compiled.selector || selector ) );\n\n\tresults = results || [];\n\n\t// Try to minimize operations if there is only one selector in the list and no seed\n\t// (the latter of which guarantees us context)\n\tif ( match.length === 1 ) {\n\n\t\t// Reduce context if the leading compound selector is an ID\n\t\ttokens = match[ 0 ] = match[ 0 ].slice( 0 );\n\t\tif ( tokens.length > 2 && ( token = tokens[ 0 ] ).type === \"ID\" &&\n\t\t\t\tcontext.nodeType === 9 && documentIsHTML &&\n\t\t\t\tjQuery.expr.relative[ tokens[ 1 ].type ] ) {\n\n\t\t\tcontext = ( jQuery.expr.find.ID(\n\t\t\t\tunescapeSelector( token.matches[ 0 ] ),\n\t\t\t\tcontext\n\t\t\t) || [] )[ 0 ];\n\t\t\tif ( !context ) {\n\t\t\t\treturn results;\n\n\t\t\t// Precompiled matchers will still verify ancestry, so step up a level\n\t\t\t} else if ( compiled ) {\n\t\t\t\tcontext = context.parentNode;\n\t\t\t}\n\n\t\t\tselector = selector.slice( tokens.shift().value.length );\n\t\t}\n\n\t\t// Fetch a seed set for right-to-left matching\n\t\ti = matchExpr.needsContext.test( selector ) ? 0 : tokens.length;\n\t\twhile ( i-- ) {\n\t\t\ttoken = tokens[ i ];\n\n\t\t\t// Abort if we hit a combinator\n\t\t\tif ( jQuery.expr.relative[ ( type = token.type ) ] ) {\n\t\t\t\tbreak;\n\t\t\t}\n\t\t\tif ( ( find = jQuery.expr.find[ type ] ) ) {\n\n\t\t\t\t// Search, expanding context for leading sibling combinators\n\t\t\t\tif ( ( seed = find(\n\t\t\t\t\tunescapeSelector( token.matches[ 0 ] ),\n\t\t\t\t\trsibling.test( tokens[ 0 ].type ) &&\n\t\t\t\t\t\ttestContext( context.parentNode ) || context\n\t\t\t\t) ) ) {\n\n\t\t\t\t\t// If seed is empty or no tokens remain, we can return early\n\t\t\t\t\ttokens.splice( i, 1 );\n\t\t\t\t\tselector = seed.length && toSelector( tokens );\n\t\t\t\t\tif ( !selector ) {\n\t\t\t\t\t\tpush.apply( results, seed );\n\t\t\t\t\t\treturn results;\n\t\t\t\t\t}\n\n\t\t\t\t\tbreak;\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\t}\n\n\t// Compile and execute a filtering function if one is not provided\n\t// Provide `match` to avoid retokenization if we modified the selector above\n\t( compiled || compile( selector, match ) )(\n\t\tseed,\n\t\tcontext,\n\t\t!documentIsHTML,\n\t\tresults,\n\t\t!context || rsibling.test( selector ) && testContext( context.parentNode ) || context\n\t);\n\treturn results;\n}\n\n// Initialize against the default document\nsetDocument();\n\njQuery.find = find;\n\n// These have always been private, but they used to be documented as part of\n// Sizzle so let's maintain them for now for backwards compatibility purposes.\nfind.compile = compile;\nfind.select = select;\nfind.setDocument = setDocument;\nfind.tokenize = tokenize;\n\nfunction dir( elem, dir, until ) {\n\tvar matched = [],\n\t\ttruncate = until !== undefined;\n\n\twhile ( ( elem = elem[ dir ] ) && elem.nodeType !== 9 ) {\n\t\tif ( elem.nodeType === 1 ) {\n\t\t\tif ( truncate && jQuery( elem ).is( until ) ) {\n\t\t\t\tbreak;\n\t\t\t}\n\t\t\tmatched.push( elem );\n\t\t}\n\t}\n\treturn matched;\n}\n\nfunction siblings( n, elem ) {\n\tvar matched = [];\n\n\tfor ( ; n; n = n.nextSibling ) {\n\t\tif ( n.nodeType === 1 && n !== elem ) {\n\t\t\tmatched.push( n );\n\t\t}\n\t}\n\n\treturn matched;\n}\n\nvar rneedsContext = jQuery.expr.match.needsContext;\n\n// rsingleTag matches a string consisting of a single HTML element with no attributes\n// and captures the element's name\nvar rsingleTag = /^<([a-z][^\\/\\0>:\\x20\\t\\r\\n\\f]*)[\\x20\\t\\r\\n\\f]*\\/?>(?:<\\/\\1>|)$/i;\n\nfunction isObviousHtml( input ) {\n\treturn input[ 0 ] === \"<\" &&\n\t\tinput[ input.length - 1 ] === \">\" &&\n\t\tinput.length >= 3;\n}\n\n// Implement the identical functionality for filter and not\nfunction winnow( elements, qualifier, not ) {\n\tif ( typeof qualifier === \"function\" ) {\n\t\treturn jQuery.grep( elements, function( elem, i ) {\n\t\t\treturn !!qualifier.call( elem, i, elem ) !== not;\n\t\t} );\n\t}\n\n\t// Single element\n\tif ( qualifier.nodeType ) {\n\t\treturn jQuery.grep( elements, function( elem ) {\n\t\t\treturn ( elem === qualifier ) !== not;\n\t\t} );\n\t}\n\n\t// Arraylike of elements (jQuery, arguments, Array)\n\tif ( typeof qualifier !== \"string\" ) {\n\t\treturn jQuery.grep( elements, function( elem ) {\n\t\t\treturn ( indexOf.call( qualifier, elem ) > -1 ) !== not;\n\t\t} );\n\t}\n\n\t// Filtered directly for both simple and complex selectors\n\treturn jQuery.filter( qualifier, elements, not );\n}\n\njQuery.filter = function( expr, elems, not ) {\n\tvar elem = elems[ 0 ];\n\n\tif ( not ) {\n\t\texpr = \":not(\" + expr + \")\";\n\t}\n\n\tif ( elems.length === 1 && elem.nodeType === 1 ) {\n\t\treturn jQuery.find.matchesSelector( elem, expr ) ? [ elem ] : [];\n\t}\n\n\treturn jQuery.find.matches( expr, jQuery.grep( elems, function( elem ) {\n\t\treturn elem.nodeType === 1;\n\t} ) );\n};\n\njQuery.fn.extend( {\n\tfind: function( selector ) {\n\t\tvar i, ret,\n\t\t\tlen = this.length,\n\t\t\tself = this;\n\n\t\tif ( typeof selector !== \"string\" ) {\n\t\t\treturn this.pushStack( jQuery( selector ).filter( function() {\n\t\t\t\tfor ( i = 0; i < len; i++ ) {\n\t\t\t\t\tif ( jQuery.contains( self[ i ], this ) ) {\n\t\t\t\t\t\treturn true;\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t} ) );\n\t\t}\n\n\t\tret = this.pushStack( [] );\n\n\t\tfor ( i = 0; i < len; i++ ) {\n\t\t\tjQuery.find( selector, self[ i ], ret );\n\t\t}\n\n\t\treturn len > 1 ? jQuery.uniqueSort( ret ) : ret;\n\t},\n\tfilter: function( selector ) {\n\t\treturn this.pushStack( winnow( this, selector || [], false ) );\n\t},\n\tnot: function( selector ) {\n\t\treturn this.pushStack( winnow( this, selector || [], true ) );\n\t},\n\tis: function( selector ) {\n\t\treturn !!winnow(\n\t\t\tthis,\n\n\t\t\t// If this is a positional/relative selector, check membership in the returned set\n\t\t\t// so $(\"p:first\").is(\"p:last\") won't return true for a doc with two \"p\".\n\t\t\ttypeof selector === \"string\" && rneedsContext.test( selector ) ?\n\t\t\t\tjQuery( selector ) :\n\t\t\t\tselector || [],\n\t\t\tfalse\n\t\t).length;\n\t}\n} );\n\n// Initialize a jQuery object\n\n// A central reference to the root jQuery(document)\nvar rootjQuery,\n\n\t// A simple way to check for HTML strings\n\t// Prioritize #id over to avoid XSS via location.hash (trac-9521)\n\t// Strict HTML recognition (trac-11290: must start with <)\n\t// Shortcut simple #id case for speed\n\trquickExpr = /^(?:\\s*(<[\\w\\W]+>)[^>]*|#([\\w-]+))$/,\n\n\tinit = jQuery.fn.init = function( selector, context ) {\n\t\tvar match, elem;\n\n\t\t// HANDLE: $(\"\"), $(null), $(undefined), $(false)\n\t\tif ( !selector ) {\n\t\t\treturn this;\n\t\t}\n\n\t\t// HANDLE: $(DOMElement)\n\t\tif ( selector.nodeType ) {\n\t\t\tthis[ 0 ] = selector;\n\t\t\tthis.length = 1;\n\t\t\treturn this;\n\n\t\t// HANDLE: $(function)\n\t\t// Shortcut for document ready\n\t\t} else if ( typeof selector === \"function\" ) {\n\t\t\treturn rootjQuery.ready !== undefined ?\n\t\t\t\trootjQuery.ready( selector ) :\n\n\t\t\t\t// Execute immediately if ready is not present\n\t\t\t\tselector( jQuery );\n\n\t\t} else {\n\n\t\t\t// Handle obvious HTML strings\n\t\t\tmatch = selector + \"\";\n\t\t\tif ( isObviousHtml( match ) ) {\n\n\t\t\t\t// Assume that strings that start and end with <> are HTML and skip\n\t\t\t\t// the regex check. This also handles browser-supported HTML wrappers\n\t\t\t\t// like TrustedHTML.\n\t\t\t\tmatch = [ null, selector, null ];\n\n\t\t\t// Handle HTML strings or selectors\n\t\t\t} else if ( typeof selector === \"string\" ) {\n\t\t\t\tmatch = rquickExpr.exec( selector );\n\t\t\t} else {\n\t\t\t\treturn jQuery.makeArray( selector, this );\n\t\t\t}\n\n\t\t\t// Match html or make sure no context is specified for #id\n\t\t\t// Note: match[1] may be a string or a TrustedHTML wrapper\n\t\t\tif ( match && ( match[ 1 ] || !context ) ) {\n\n\t\t\t\t// HANDLE: $(html) -> $(array)\n\t\t\t\tif ( match[ 1 ] ) {\n\t\t\t\t\tcontext = context instanceof jQuery ? context[ 0 ] : context;\n\n\t\t\t\t\t// Option to run scripts is true for back-compat\n\t\t\t\t\t// Intentionally let the error be thrown if parseHTML is not present\n\t\t\t\t\tjQuery.merge( this, jQuery.parseHTML(\n\t\t\t\t\t\tmatch[ 1 ],\n\t\t\t\t\t\tcontext && context.nodeType ? context.ownerDocument || context : document$1,\n\t\t\t\t\t\ttrue\n\t\t\t\t\t) );\n\n\t\t\t\t\t// HANDLE: $(html, props)\n\t\t\t\t\tif ( rsingleTag.test( match[ 1 ] ) && jQuery.isPlainObject( context ) ) {\n\t\t\t\t\t\tfor ( match in context ) {\n\n\t\t\t\t\t\t\t// Properties of context are called as methods if possible\n\t\t\t\t\t\t\tif ( typeof this[ match ] === \"function\" ) {\n\t\t\t\t\t\t\t\tthis[ match ]( context[ match ] );\n\n\t\t\t\t\t\t\t// ...and otherwise set as attributes\n\t\t\t\t\t\t\t} else {\n\t\t\t\t\t\t\t\tthis.attr( match, context[ match ] );\n\t\t\t\t\t\t\t}\n\t\t\t\t\t\t}\n\t\t\t\t\t}\n\n\t\t\t\t\treturn this;\n\n\t\t\t\t// HANDLE: $(#id)\n\t\t\t\t} else {\n\t\t\t\t\telem = document$1.getElementById( match[ 2 ] );\n\n\t\t\t\t\tif ( elem ) {\n\n\t\t\t\t\t\t// Inject the element directly into the jQuery object\n\t\t\t\t\t\tthis[ 0 ] = elem;\n\t\t\t\t\t\tthis.length = 1;\n\t\t\t\t\t}\n\t\t\t\t\treturn this;\n\t\t\t\t}\n\n\t\t\t// HANDLE: $(expr) & $(expr, $(...))\n\t\t\t} else if ( !context || context.jquery ) {\n\t\t\t\treturn ( context || rootjQuery ).find( selector );\n\n\t\t\t// HANDLE: $(expr, context)\n\t\t\t// (which is just equivalent to: $(context).find(expr)\n\t\t\t} else {\n\t\t\t\treturn this.constructor( context ).find( selector );\n\t\t\t}\n\t\t}\n\n\t};\n\n// Give the init function the jQuery prototype for later instantiation\ninit.prototype = jQuery.fn;\n\n// Initialize central reference\nrootjQuery = jQuery( document$1 );\n\nvar rparentsprev = /^(?:parents|prev(?:Until|All))/,\n\n\t// Methods guaranteed to produce a unique set when starting from a unique set\n\tguaranteedUnique = {\n\t\tchildren: true,\n\t\tcontents: true,\n\t\tnext: true,\n\t\tprev: true\n\t};\n\njQuery.fn.extend( {\n\thas: function( target ) {\n\t\tvar targets = jQuery( target, this ),\n\t\t\tl = targets.length;\n\n\t\treturn this.filter( function() {\n\t\t\tvar i = 0;\n\t\t\tfor ( ; i < l; i++ ) {\n\t\t\t\tif ( jQuery.contains( this, targets[ i ] ) ) {\n\t\t\t\t\treturn true;\n\t\t\t\t}\n\t\t\t}\n\t\t} );\n\t},\n\n\tclosest: function( selectors, context ) {\n\t\tvar cur,\n\t\t\ti = 0,\n\t\t\tl = this.length,\n\t\t\tmatched = [],\n\t\t\ttargets = typeof selectors !== \"string\" && jQuery( selectors );\n\n\t\t// Positional selectors never match, since there's no _selection_ context\n\t\tif ( !rneedsContext.test( selectors ) ) {\n\t\t\tfor ( ; i < l; i++ ) {\n\t\t\t\tfor ( cur = this[ i ]; cur && cur !== context; cur = cur.parentNode ) {\n\n\t\t\t\t\t// Always skip document fragments\n\t\t\t\t\tif ( cur.nodeType < 11 && ( targets ?\n\t\t\t\t\t\ttargets.index( cur ) > -1 :\n\n\t\t\t\t\t\t// Don't pass non-elements to jQuery#find\n\t\t\t\t\t\tcur.nodeType === 1 &&\n\t\t\t\t\t\t\tjQuery.find.matchesSelector( cur, selectors ) ) ) {\n\n\t\t\t\t\t\tmatched.push( cur );\n\t\t\t\t\t\tbreak;\n\t\t\t\t\t}\n\t\t\t\t}\n\t\t\t}\n\t\t}\n\n\t\treturn this.pushStack( matched.length > 1 ? jQuery.uniqueSort( matched ) : matched );\n\t},\n\n\t// Determine the position of an element within the set\n\tindex: function( elem ) {\n\n\t\t// No argument, return index in parent\n\t\tif ( !elem ) {\n\t\t\treturn ( this[ 0 ] && this[ 0 ].parentNode ) ? this.first().prevAll().length : -1;\n\t\t}\n\n\t\t// Index in selector\n\t\tif ( typeof elem === \"string\" ) {\n\t\t\treturn indexOf.call( jQuery( elem ), this[ 0 ] );\n\t\t}\n\n\t\t// Locate the position of the desired element\n\t\treturn indexOf.call( this,\n\n\t\t\t// If it receives a jQuery object, the first element is used\n\t\t\telem.jquery ? elem[ 0 ] : elem\n\t\t);\n\t},\n\n\tadd: function( selector, context ) {\n\t\treturn this.pushStack(\n\t\t\tjQuery.uniqueSort(\n\t\t\t\tjQuery.merge( this.get(), jQuery( selector, context ) )\n\t\t\t)\n\t\t);\n\t},\n\n\taddBack: function( selector ) {\n\t\treturn this.add( selector == null ?\n\t\t\tthis.prevObject : this.prevObject.filter( selector )\n\t\t);\n\t}\n} );\n\nfunction sibling( cur, dir ) {\n\twhile ( ( cur = cur[ dir ] ) && cur.nodeType !== 1 ) {}\n\treturn cur;\n}\n\njQuery.each( {\n\tparent: function( elem ) {\n\t\tvar parent = elem.parentNode;\n\t\treturn parent && parent.nodeType !== 11 ? parent : null;\n\t},\n\tparents: function( elem ) {\n\t\treturn dir( elem, \"parentNode\" );\n\t},\n\tparentsUntil: function( elem, _i, until ) {\n\t\treturn dir( elem, \"parentNode\", until );\n\t},\n\tnext: function( elem ) {\n\t\treturn sibling( elem, \"nextSibling\" );\n\t},\n\tprev: function( elem ) {\n\t\treturn sibling( elem, \"previousSibling\" );\n\t},\n\tnextAll: function( elem ) {\n\t\treturn dir( elem, \"nextSibling\" );\n\t},\n\tprevAll: function( elem ) {\n\t\treturn dir( elem, \"previousSibling\" );\n\t},\n\tnextUntil: function( elem, _i, until ) {\n\t\treturn dir( elem, \"nextSibling\", until );\n\t},\n\tprevUntil: function( elem, _i, until ) {\n\t\treturn dir( elem, \"previousSibling\", until );\n\t},\n\tsiblings: function( elem ) {\n\t\treturn siblings( ( elem.parentNode || {} ).firstChild, elem );\n\t},\n\tchildren: function( elem ) {\n\t\treturn siblings( elem.firstChild );\n\t},\n\tcontents: function( elem ) {\n\t\tif ( elem.contentDocument != null &&\n\n\t\t\t// Support: IE 11+\n\t\t\t//