Files
olivv-cs-ComfyUI-FunPack/modules/loaders/diffusion_model/nodes.py
T
digital-garbageandClaude Opus 5 20a4c39db3 loaders: compute_dtype cast the input and left the weights behind
set_model_compute_dtype sets force_cast_weights on purpose: it is what
casts each layer's weights to match. Clearing it afterwards left the model
casting its INPUT to the chosen dtype while the weights stayed as loaded,
so the first Linear of the first sampling step raised "mat1 and mat2 must
have the same dtype".

Picking any compute dtype other than default therefore broke sampling on a
full-VRAM load. It survives on constrained hardware only because the
low-VRAM path forces the cast back on by itself, which is why nobody hit it.

Straight port from v4, where the same line has shipped since 3837f8f. The
branch had no test at all; it now has one that fails against the old line.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-29 12:06:45 +03:00

94 lines
4.7 KiB
Python

"""The diffusion model loader.
Picking a file and how to read it. Every option here changes what the loaded
model IS, which is why they are node widgets and not settings: nobody downstream
can supply them, and a saved workflow has to record them.
"""
import comfy.sd
import comfy.utils
import folder_paths
from comfy_api.latest import io
from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices,
attention_override, dtype_of, set_fp16_accumulation,
weight_model_options)
class FunPackDiffusionModelLoader(io.ComfyNode):
@classmethod
def define_schema(cls) -> io.Schema:
return io.Schema(
node_id="FunPackDiffusionModelLoader",
display_name="FunPack Diffusion Model Loader",
category="FunPack/Loaders",
description="Load a diffusion model, choosing precision and attention backend.",
inputs=[
io.Combo.Input("model_name", options=folder_paths.get_filename_list("diffusion_models")),
io.Combo.Input("weight_dtype", options=WEIGHT_DTYPES, default="default",
tooltip="How the weights are stored. 'default' keeps whatever "
"the file already is."),
io.Combo.Input("compute_dtype", options=COMPUTE_DTYPES, default="default",
tooltip="What the maths runs in, independently of storage."),
io.Combo.Input("attention", options=attention_choices(), default="default",
tooltip="Attention backend. Only what this machine can run is "
"listed; 'default' keeps ComfyUI's launch choice."),
io.Boolean.Input("fp16_accumulation", default=False, optional=True,
tooltip="Faster fp16 matmuls where the torch build supports it."),
],
outputs=[
io.Model.Output(display_name="model"),
io.String.Output(display_name="status"),
],
)
@classmethod
def execute(cls, model_name: str, weight_dtype: str, compute_dtype: str,
attention: str, fp16_accumulation: bool = False) -> io.NodeOutput:
notes = [f"FunPack Diffusion Model Loader | {model_name}"]
accumulation = set_fp16_accumulation(fp16_accumulation)
if fp16_accumulation and accumulation is None:
notes.append("fp16_accumulation: unsupported by this torch build, ignored")
elif accumulation is not None:
notes.append(f"fp16_accumulation: {accumulation}")
model_options = weight_model_options(weight_dtype)
notes.append(f"weight dtype: {weight_dtype}")
path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
state_dict, metadata = comfy.utils.load_torch_file(path, return_metadata=True)
model = comfy.sd.load_diffusion_model_state_dict(
state_dict, model_options=model_options, metadata=metadata)
if model is None:
# comfy returns None rather than raising, and a None reaching the
# sampler fails much further away than the mistake.
raise RuntimeError(
f"Could not detect a diffusion model in {model_name}. Loading a text "
f"encoder or a VAE file here is the usual cause.")
dtype = dtype_of(compute_dtype)
if dtype is not None and hasattr(model, "set_model_compute_dtype"):
# Do NOT clear force_cast_weights afterwards. set_model_compute_dtype
# sets it deliberately, and it is what casts each layer's WEIGHTS to
# match. Without it the model still casts its INPUT to the requested
# dtype (_apply_model reads manual_cast_dtype), so the first Linear
# gets a bf16 activation against fp32 weights and sampling dies with
# "mat1 and mat2 must have the same dtype". v4 cleared it and shipped
# that way; it survives on constrained hardware only because the
# low-VRAM path forces the cast back on regardless.
model.set_model_compute_dtype(dtype)
notes.append(f"compute dtype: {compute_dtype}")
elif dtype is not None:
notes.append("compute dtype: unsupported by this ComfyUI, ignored")
override = attention_override(attention)
if override is not None:
model.model_options.setdefault("transformer_options", {})[
"optimized_attention_override"] = override
notes.append(f"attention: {attention}")
else:
notes.append("attention: default (as launched)")
return io.NodeOutput(model, "\n".join(notes))