olivv-cs-ComfyUI-FunPack/modules/loaders/common.py

"""Shared loading options.

Not a module: it announces nothing and ships no node. `registry.scan` only walks
two levels of DIRECTORIES, so a plain file at the domain level is a helper the
loaders import and nothing else.

These options are on the loader node rather than in a settings payload, and the
line is worth stating: a loader widget describes the FILE and how it is read,
which the loader cannot be handed by anyone else. Feature knobs -- strengths,
sigma windows, anything a modifier owns -- are settings, and never appear here.
"""

import torch

from .._core import log

# Named exactly as core's UNETLoader names them where they overlap, so a value
# copied from a stock workflow means the same thing here.
WEIGHT_DTYPES = ["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"]
COMPUTE_DTYPES = ["default", "fp16", "bf16", "fp32"]
VAE_DTYPES = ["bfloat16", "float16", "float32"]

_BY_NAME = {
    "fp16": torch.float16,
    "bf16": torch.bfloat16,
    "fp32": torch.float32,
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
    "fp8_e4m3fn": getattr(torch, "float8_e4m3fn", None),
    "fp8_e5m2": getattr(torch, "float8_e5m2", None),
    # fp8_e4m3fn plus comfy's fp8 matmul path; see weight_model_options.
    "fp8_e4m3fn_fast": getattr(torch, "float8_e4m3fn", None),
}


def dtype_of(name):
    """The torch dtype a choice names, or None for 'default'/an unsupported build."""
    return _BY_NAME.get(str(name or "default"))


def unsupported_dtypes(model, weight_dtype, compute_dtype):
    """A forced dtype the model does not run in, one line each. A model says which it runs in (H3: bf16/fp32;
    v4 measured fp16 breaking it). Allowed, but said."""
    supported = getattr(getattr(getattr(model, "model", None), "model_config", None), "supported_inference_dtypes", None) or []
    runs_in = ", ".join(str(d).replace("torch.", "") for d in supported)
    return [f"{what} dtype {picked} is not one this model runs in ({runs_in})"
            for what, picked in (("weight", weight_dtype), ("compute", compute_dtype))
            if picked in ("fp16", "bf16", "fp32") and supported and dtype_of(picked) not in supported]


_SHORT = {"bfloat16": "bf16", "float16": "fp16", "float32": "fp32"}


def stored_weights(state_dict):
    """What a file's weights are stored as, by element count: "bf16", or "int8+bf16" for a
    mixed file. Read from the tensors, not the weight_dtype widget: a pre-quantized file
    loaded with weight_dtype bf16 still runs int8 layers, and the log must say so."""
    sizes = {}
    for t in state_dict.values():
        if torch.is_tensor(t) and t.dim() >= 2:
            name = str(t.dtype).replace("torch.", "")
            sizes[_SHORT.get(name, name)] = sizes.get(_SHORT.get(name, name), 0) + t.numel()
    total = sum(sizes.values())
    return "+".join(n for n in sorted(sizes, key=sizes.get, reverse=True) if sizes[n] > total * 0.1) or None


def slow_int8_build():
    """Why int8 models sample slowly on this torch build, or None.

    Measured on an RTX PRO 6000 Blackwell, same 4-step H3 run: torch 2.10+cu128 sampled the
    int8 model ~2x SLOWER than bf16; the same torch from the cu130 index made int8 ~25% faster.
    """
    cuda = getattr(torch.version, "cuda", None)
    try:
        if not cuda or not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10:
            return None
        if int(cuda.split(".")[0]) >= 13:
            return None
    except (RuntimeError, ValueError):
        return None
    # --force-reinstall is required: pip counts 2.10.0+cu128 as already satisfying torch==2.10.0.
    return (f"torch is built for CUDA {cuda}: on this Blackwell GPU int8 models sample about 2x slower "
            f"than on a CUDA 13 build. Stop ComfyUI and reinstall torch, torchvision and torchaudio "
            f"(same versions) with --force-reinstall --index-url https://download.pytorch.org/whl/cu130")


def weight_model_options(weight_dtype):
    """comfy `model_options` for a weight dtype. The mapping core's UNETLoader uses."""
    options = {}
    dtype = dtype_of(weight_dtype)
    if dtype is not None:
        options["dtype"] = dtype
    if weight_dtype == "fp8_e4m3fn_fast":
        options["fp8_optimizations"] = True
    return options


def attention_choices():
    """Backends this ComfyUI can actually run.

    Read from ComfyUI's own registry rather than a private import, so the list
    offers exactly what the machine has -- sage where SageAttention is installed,
    flash where flash-attn is -- and never lies about the rest. 'default' leaves
    whatever ComfyUI itself selected from its launch flags.
    """
    names = []
    try:
        from comfy.ldm.modules.attention import REGISTERED_ATTENTION_FUNCTIONS
        names = sorted(REGISTERED_ATTENTION_FUNCTIONS.keys())
    except Exception:                            # noqa: BLE001 -- older ComfyUI
        pass
    return ["default"] + names


def attention_override(name):
    """A transformer_options override routing every attention call to `name`.

    ComfyUI wraps each implementation with `wrap_attn`, which hands the override
    the original function plus the call's arguments -- so calling the WRAPPED
    replacement would re-enter that machinery. The unwrapped one is what gets
    called.
    """
    if not name or name == "default":
        return None
    try:
        from comfy.ldm.modules.attention import get_attention_function
    except ImportError:
        log.warning("attention", "this ComfyUI has no attention registry, so the "
                                 "chosen backend cannot be applied")
        return None
    chosen = get_attention_function(name, None)
    if chosen is None:
        log.warning("attention", f"backend {name!r} is not available on this machine, "
                                 f"so it was not applied")
        return None
    inner = getattr(chosen, "__wrapped__", chosen)

    def override(_func, *args, **kwargs):
        return inner(*args, **kwargs)

    # Some backends carry a zero-copy path that takes the attention tensors still
    # in their containers, and `wrap_attn` only offers it when the OVERRIDE has
    # the attribute -- otherwise it unpacks each container and calls the slow way.
    # Forgetting to forward it is invisible: the backend is selected, the maths is
    # right, and the fast path it exists for is silently never used. ComfyUI's own
    # `ModelPatcher.set_model_optimized_attention` forwards it for this reason.
    container_function = getattr(chosen, "container_function", None)
    if container_function is not None:
        override.container_function = container_function

    return override


def set_fp16_accumulation(enabled):
    """torch's fp16 accumulation switch, or None where the build has no such knob.

    PROCESS-WIDE, and there is no scoped version of it: torch has one flag for the
    whole interpreter. So this cannot mean "for this model" however it is spelled
    on a node -- a second loader setting it differently changes the first model's
    maths, and because ComfyUI caches loader outputs the last value set outlives
    the run that set it, reaching models loaded by core's own nodes.

    It is therefore never changed silently. Every transition says so once, which
    is the difference between a documented global and an invisible one. The real
    fix is to move this to a machine-level setting, where its scope is honest.
    """
    matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
    if matmul is None or not hasattr(matmul, "allow_fp16_accumulation"):
        return None
    wanted = bool(enabled)
    before = bool(getattr(matmul, "allow_fp16_accumulation", False))
    matmul.allow_fp16_accumulation = wanted
    if before != wanted:
        log.alert("fp16 accumulation",
                  f"now {wanted} for EVERY model in this process, not just this one -- "
                  f"torch has a single process-wide flag")
    return wanted