Ported from v4, minus two mechanisms that are no longer needed. The LoRA loader takes one LoRA and chains, so the FlexibleOptionalInputType dict subclass that lied to ComfyUI's validator has nothing left to do. The CLIP loader uses numbered slots like core's DualCLIPLoader instead of a JSON array in a string widget. Autogrow looked like the right answer for the encoder list and is not: it sets force_input on any widget input, so a list of combos renders as sockets and the file picker stops being a picker. Every core use of it is a connection type. GGUF and SLA are deliberately not here. Both are capabilities rather than loader options, and both belong in modules of their own. The tests follow what the loaders PRODUCE rather than stopping at the call: the attention override is carried into comfy's real dispatch and checked against every backend this machine registers, because an override that only fails at the first attention call fails at sampling, not at load. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
87 lines
4.1 KiB
Python
87 lines
4.1 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"):
|
|
model.set_model_compute_dtype(dtype)
|
|
model.force_cast_weights = False
|
|
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))
|