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>
94 lines
4.7 KiB
Python
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))
|