82 lines
2.5 KiB
Python
82 lines
2.5 KiB
Python
import comfy
|
|
import folder_paths
|
|
import torch
|
|
|
|
from dfloat11 import DFloat11Model
|
|
|
|
class DFloat11ModelLoader:
|
|
"""
|
|
A custom node to load a DFloat11 diffusion model from the `diffusion_models` directory.
|
|
|
|
DFloat11 models are >30% smaller than their float16 counterparts, yet produce bit-for-bit identical outputs.
|
|
"""
|
|
|
|
MODEL_TO_PATTERN_DICT = {
|
|
"Flux": {
|
|
"double_blocks\.\d+": (
|
|
"img_mod.lin",
|
|
"img_attn.qkv",
|
|
"img_attn.proj",
|
|
"img_mlp.0",
|
|
"img_mlp.2",
|
|
"txt_mod.lin",
|
|
"txt_attn.qkv",
|
|
"txt_attn.proj",
|
|
"txt_mlp.0",
|
|
"txt_mlp.2",
|
|
),
|
|
"single_blocks\.\d+": (
|
|
"linear1",
|
|
"linear2",
|
|
"modulation.lin",
|
|
),
|
|
},
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"dfloat11_model_name": (folder_paths.get_filename_list("diffusion_models"),),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL",)
|
|
FUNCTION = "load_dfloat11_model"
|
|
CATEGORY = "DFloat11"
|
|
|
|
def load_dfloat11_model(self, dfloat11_model_name):
|
|
dfloat11_model_path = folder_paths.get_full_path_or_raise("diffusion_models", dfloat11_model_name)
|
|
sd = comfy.utils.load_torch_file(dfloat11_model_path)
|
|
|
|
if not any(k.endswith("encoded_exponent") for k in sd.keys()):
|
|
raise ValueError(f"The model '{dfloat11_model_name}' is not a DFloat11 model.")
|
|
|
|
load_device = comfy.model_management.get_torch_device()
|
|
offload_device = comfy.model_management.unet_offload_device()
|
|
|
|
model_config = comfy.sd.model_detection.model_config_from_unet(sd, "")
|
|
model_config.set_inference_dtype(torch.bfloat16, torch.bfloat16)
|
|
model = model_config.get_model(sd, "")
|
|
model = model.to(offload_device)
|
|
|
|
DFloat11Model.from_single_file(
|
|
dfloat11_model_path,
|
|
pattern_dict=DFloat11ModelLoader.MODEL_TO_PATTERN_DICT[type(model_config).__name__],
|
|
bfloat16_model=model.diffusion_model,
|
|
device=offload_device,
|
|
)
|
|
|
|
return (
|
|
comfy.model_patcher.ModelPatcher(model, load_device=load_device, offload_device=offload_device),
|
|
)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DFloat11ModelLoader": DFloat11ModelLoader,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DFloat11ModelLoader": "DFloat11 Model Loader",
|
|
}
|