Merge pull request #5011 from vladmandic/dev

dev merge

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.

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

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,

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 ({}, {}, [])

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",

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"
   }
 }

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"
   }
 }

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"
   }
 }

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"
   }

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')

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)
 

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"])

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

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")

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
 

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)

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

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__}')

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 []

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 = {}

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'],

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:

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)}')

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

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):

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}")

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 '<frozen' not in x
              and '{built-in' not in x
@@ -92,6 +95,23 @@ def profile(profiler, msg: str, n: int = 16):
     log.debug(f'Profile {msg}: {txt}')
 
 
+def profile_start():
+    global _profiler # pylint: disable=global-statement
+    if _profiler is not None:
+        _profiler.disable()
+        _profiler.clear()
+        _profiler.enable()
+    else:
+        import cProfile
+        _profiler = cProfile.Profile()
+        _profiler.enable()
+
+
+def profile_stop():
+    if _profiler is not None:
+        _profiler.disable()
+
+
 def profile_torch(profiler, msg: str):
     profiler.stop()
     lines = profiler.key_averages().table(sort_by="cpu_time_total", row_limit=12)

modules/files_cache.py

@@ -1,13 +1,13 @@
 from typing import Union
-import itertools
 import os
 from collections import UserDict
 from collections.abc import Callable, Iterator
 from dataclasses import dataclass, field
+from functools import lru_cache
 from modules.logger import log
 
-
 do_cache_folders = os.environ.get('SD_NO_CACHE', None) is None
+
 class Directory: # forward declaration
     ...
 
@@ -20,34 +20,34 @@ DirectoryIterator = Iterator[Directory]
 DirectoryCollection = dict[str, Directory]
 ExtensionFilter = Callable
 ExtensionList = list[str]
-RecursiveType = Union[bool,Callable]
+RecursiveType = Union[bool, Callable]
 
 
-def real_path(directory_path:str) -> 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({})

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',
     )

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)

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)
 

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)

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:

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.

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',
     )

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

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

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)}')

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'

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):

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
 
 

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

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)}')
 
 

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]

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,

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

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

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

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)

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

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()

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:

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

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:

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):

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

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)
 
 

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)

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:

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):

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

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)

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)

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

modules/seedvr/src/common/decorators.py

@@ -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

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):

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)

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
 

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)
 

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)

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,

modules/shared_items.py

@@ -84,6 +84,8 @@ pipelines = {
     'XOmni': None,
     'ZetaChroma': None,
     'Boogu': None,
+    'SeFi': None,
+    'MageFlow': None,
 }
 
 

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

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)
 

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

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")
 

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"<div class='performance hint' id='control-performance'><p>{elapsed_text} {summary} {gpu} {cpu}</p></div>"
 
@@ -247,7 +248,7 @@ def create_ui(_blocks: gr.Blocks=None):
                     gr.HTML('<span id="control-output-button">Output</p>')
                     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'):

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("<h2>Balanced Offload</h2>", "", 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 }),

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)

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
 

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))) + '<br>' + 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('<br>')[-1]

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

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]

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]

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]

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'):

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')

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

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

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

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

modules/video_models/video_prompt.py

@@ -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

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

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))

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',
     )

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,
                 )