Claude
6ae2247e78
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjVsBg3ixjnT1Fgdy12
261 lines
12 KiB
Python
261 lines
12 KiB
Python
"""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 json
|
|
|
|
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]
|
|
|
|
|
|
def fp32_over_bf16(weight_dtype, compute_dtype, stored=None):
|
|
"""Why fp32 compute over bf16 weights is noisy, or None. It runs, but the activations go fp32 while the norm layers keep
|
|
their bf16 weights, so torch warns "Mismatch dtype between input and weight". Allowed, but said."""
|
|
if compute_dtype == "fp32" and (weight_dtype == "bf16" or (weight_dtype == "default" and stored == "bf16")):
|
|
return "bf16 weights dtype with fp32 precision is a bad combo for H3. Don't expect miracles: it runs, but slowly."
|
|
|
|
|
|
_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
|
|
|
|
|
|
def head_bank_metadata(state_dict, metadata):
|
|
"""File metadata that declares how many output heads a few-step H3 file stacks, or the input untouched.
|
|
|
|
A file with several stacked heads per stream normally says so in its safetensors
|
|
config (`transformer.pdd_steps`). A re-saved copy can lose that line while keeping
|
|
explicit head sizes, and ComfyUI then builds a single head and fails on the biases.
|
|
The count is in the weights anyway: the patch projections hold the true per-head
|
|
width, and the head bias holds that width times the count.
|
|
|
|
Reads shapes only, so it works on any device, dtype or quantised storage.
|
|
"""
|
|
try:
|
|
bias = state_dict["final_layer.video_out.bias"]
|
|
abias = state_dict["final_layer.audio_out.bias"]
|
|
video = state_dict["video_patch_proj.weight"].shape[1]
|
|
audio = state_dict["audio_patch_proj.weight"].shape[1]
|
|
except (KeyError, AttributeError, IndexError):
|
|
return metadata, None
|
|
if not video or not audio or bias.shape[0] % video or abias.shape[0] % audio:
|
|
return metadata, None
|
|
steps = bias.shape[0] // video
|
|
if steps < 2 or abias.shape[0] // audio != steps:
|
|
return metadata, None
|
|
try:
|
|
config = json.loads((metadata or {}).get("config") or "{}")
|
|
except ValueError:
|
|
config = {}
|
|
transformer = config.setdefault("transformer", {})
|
|
if transformer.get("pdd_steps", 1) == steps:
|
|
return metadata, None
|
|
declared = transformer.get("pdd_steps")
|
|
transformer["pdd_steps"] = steps
|
|
fixed = dict(metadata or {})
|
|
fixed["config"] = json.dumps(config, separators=(",", ":"))
|
|
return fixed, (f"output heads: the file stacks {steps} per stream but its metadata "
|
|
f"{'says ' + str(declared) if declared else 'does not say so'}; using {steps}")
|
|
|
|
|
|
_loaded_as = {} # model file -> [run id, the settings seen in that run, the settings seen in the run before]
|
|
|
|
|
|
def free_stale_copy_after_run(name, settings) -> bool:
|
|
"""Loading a file again with other settings (attention, SLA method, dtype...) leaves the earlier copy in ComfyUI's node cache,
|
|
until RAM runs short and ComfyUI evicts it half way, which is when "memory leak with model X" shows up. So when the same file
|
|
comes back with settings the run before did not have, ask ComfyUI to unload its models and reset its cache once this run ends
|
|
(its own /free flag): the next run loads the file once more, and only one copy was ever around. Two loaders of one file in one
|
|
graph are no change: each run sees the same set of settings, so nothing is flagged and nothing reloads every run."""
|
|
try:
|
|
from server import PromptServer
|
|
server = PromptServer.instance
|
|
running = list((getattr(server.prompt_queue, "currently_running", None) or {}).values()) # ComfyUI's running item: (number, prompt id, ...)
|
|
run = running[0][1] if running else object() # no run id (a plain script, a test): every call is its own run
|
|
except Exception: # noqa: BLE001 -- no server (tests, a plain script): nothing to ask
|
|
return False
|
|
key = json.dumps(settings, sort_keys=True, default=str)
|
|
seen = _loaded_as.get(name)
|
|
if seen is None or seen[0] != run:
|
|
seen = _loaded_as[name] = [run, set(), seen[1] if seen else None]
|
|
seen[1].add(key)
|
|
previous = seen[2]
|
|
if previous is None or key in previous:
|
|
return False
|
|
try:
|
|
server.prompt_queue.set_flag("free_memory", True)
|
|
except Exception as exc: # noqa: BLE001 -- a ComfyUI without the flag: said, so the leak warning is not a mystery
|
|
log.warning("FunUI Diffusion Model Loader", f"{name} came back with other settings, but this ComfyUI cannot be asked to free the earlier copy ({type(exc).__name__}); "
|
|
"a 'memory leak with model' warning may follow")
|
|
return False
|
|
log.info("FunUI Diffusion Model Loader",
|
|
f"{name} was loaded before with other settings: the earlier copy is freed when this run ends, so the next run loads it once more")
|
|
return True
|