olivv-cs-ComfyUI-FunPack/modules/loaders/diffusion_model/nodes.py

"""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 ..._core import log
from .. import gguf_support
from ..common import (COMPUTE_DTYPES, WEIGHT_DTYPES, attention_choices,
                      attention_override, dtype_of, set_fp16_accumulation,
                      weight_model_options)


def _model_file_choices():
    """Diffusion model files, `.gguf` included.

    Core's extension set has no `.gguf`, so those files are on disk and
    invisible to every picker. Appended rather than merged in sorted order, so
    an existing pipeline's saved choice keeps its position in the list and
    nothing a user already picked moves.
    """
    return list(folder_paths.get_filename_list("diffusion_models")) + \
        gguf_support.gguf_names("diffusion_models")


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=_model_file_choices()),
                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}")

        if gguf_support.is_gguf(model_name):
            path = gguf_support.gguf_path("diffusion_models", model_name)
            if not path:
                raise RuntimeError(f"{model_name} is no longer where it was listed from.")
            misnamed = False
        else:
            path = folder_paths.get_full_path_or_raise("diffusion_models", model_name)
            # The CONTENT decides, not the extension. A .gguf renamed to
            # .safetensors used to reach the safetensors parser and fail with
            # a UTF-8 decode error from reading a binary header as JSON --
            # true, and no help at all.
            misnamed = gguf_support.has_gguf_magic(path)
        if gguf_support.is_gguf(model_name) or misnamed:
            if misnamed:
                notes.append(f"{model_name} is named .safetensors but is a GGUF "
                             f"container -- loaded as GGUF")
            state_dict, gguf_options, gguf_note = gguf_support.load_state_dict(path)
            # The quantized path needs its own torch operations, and they must
            # not be lost to the dtype options above -- a GGUF loaded with
            # stock ops would try to matmul block-quantized storage.
            model_options = {**model_options, **gguf_options}
            metadata = None
            notes.append(gguf_note)
        else:
            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)")

        # Never reach into the model's shape to describe it: a log line that
        # assumes structure can fail the load it was only meant to narrate.
        kind = type(getattr(model, "model", model)).__name__
        log.info("FunPack Diffusion Model Loader",
                 f"{model_name} loaded as {kind}, weights "
                 f"{weight_dtype}, compute {compute_dtype}, attention "
                 f"{attention if override is not None else 'default (as launched)'}")
        return io.NodeOutput(model, "\n".join(notes))