CheckpointLoaderKJ

The CheckpointLoaderKJ node builds upon the standard loading process by adding control over various weight precisions (such as FP8, FP16, BF16) and computational precision.

CheckpointLoaderKJ is an experimental model loading node designed to load Stable Diffusion models from checkpoint files. Its core functionality includes loading the model, CLIP text encoder, and VAE decoder, while supporting specific performance optimizations and precision adjustments for the model.

Node Functionality

This node extends the standard loading process by adding control over various weight precisions (e.g., FP8, FP16, BF16) and computational precision. Its key technical mechanisms include optionally using CublasLinear to patch PyTorch's linear layers for improved computational efficiency, and integrating SageAttention to replace the default attention mechanism. These optimizations aim to enhance inference speed or reduce VRAM usage for specific hardware (such as GPUs supporting FP8).

Node Parameter Description - CheckpointLoaderKJ

Connection Inputs (Inputs)

Parameter Name	Data Type	Required	Default Value	Value Range/Options	Description
`ckpt_name`	FOLDER_PATHS	Yes	-	-	ckpt_name (Input Parameter): Specifies the name of the checkpoint (model) file to load. This file must be located within the model directory specified by ComfyUI.

Control Parameters (Parameters)

Parameter Name	Data Type	Required	Default Value	Value Range/Options	Description
`weight_dtype`	COMBO	Yes	default	"default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2", "fp16", "bf16", "fp32"	Controls memory usage and computational precision.
`compute_dtype`	COMBO	Yes	default	"default", "fp16", "bf16", "fp32"	fp16, bf16, fp32. Affects numerical precision during the forward propagation process.
`patch_cublaslinear`	BOOLEAN	Yes	False	-	Enable or disable the cublas_ops arg
`sage_attention`	COMBO	Yes	False	false, true	Patch comfy attention to use sageattn.
`enable_fp16_accumulation`	BOOLEAN	Yes	False	-	Enable torch.backends.cuda.matmul.allow_fp16_accumulation, requires minimum PyTorch version 2.7.1

Output

Parameter Name	Data Type	Description
MODEL	MODEL	MODEL (Output Parameter): The loaded Stable Diffusion main model, used for image denoising and generation.
CLIP	CLIP	CLIP (Output Parameter): The CLIP text encoder, used to convert text prompts into embedding vectors understandable by the model.
VAE	VAE	VAE (Output Parameter): The variational autoencoder, responsible for encoding and decoding latent space representations into pixel images.

Usage Scenarios

This node is particularly useful in workflows requiring precise control over model precision to balance speed and quality. For example, when deploying optimized models for batch image generation, users can first load the model with FP8 precision using this node, then connect it to a KSampler for efficient inference. It can also be used to test the impact of different attention mechanisms (such as SageAttention) on generation results.

Notes

Please note that this is a node marked as "experimental"; its functionality may be unstable or change with version updates. Enabling features like enable_fp16_accumulation requires specific versions of PyTorch (2.7.1 or higher), and some optimizations (like FP8) require hardware support.