From ff143a18969e1fa6f25d8d7cf31b25e9e561faa8 Mon Sep 17 00:00:00 2001 From: Tianyi Zhang Date: Tue, 26 Aug 2025 03:00:09 -0500 Subject: [PATCH] Add files --- README.md | 70 ++++++++++++++++++++++++++++++++++ __init__.py | 3 ++ dfloat11_model_loader.py | 81 ++++++++++++++++++++++++++++++++++++++++ requirements.txt | 1 + 4 files changed, 155 insertions(+) create mode 100644 README.md create mode 100644 __init__.py create mode 100644 dfloat11_model_loader.py create mode 100644 requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..5f3e2f7 --- /dev/null +++ b/README.md @@ -0,0 +1,70 @@ +# ComfyUI Plugin for DFloat11 + +[![PyPI Downloads](https://static.pepy.tech/personalized-badge/dfloat11?period=total\&units=INTERNATIONAL_SYSTEM\&left_color=BLACK\&right_color=GREEN\&left_text=downloads)](https://pepy.tech/projects/dfloat11) +[![arXiv](https://img.shields.io/badge/arXiv-2504.11651-b31b1b.svg)](https://arxiv.org/abs/2504.11651) +[![Hugging Face](https://img.shields.io/badge/Model-%F0%9F%A4%97-yellow.svg)](https://huggingface.co/DFloat11) + +This repository provides the **ComfyUI plugin for DFloat11 models**. + +DFloat11 reduces model size by more than **30%** while producing **bit-for-bit identical outputs** to the original. Unlike quantization techniques which trade quality for size, DFloat11 is a **lossless compression method**, preserving model output quality fully while supporting efficient inference. + +Currently, only **FLUX.1 models** are supported. Support for additional models is planned. Please feel free to [open an issue](https://github.com/LeanModels/ComfyUI-DFloat11/issues) and let us know which ones you'd like to see next. + +--- + +## Features + +* 🚀 **Fully Lossless** – 100% identical outputs to the original model +* 📦 **>30% smaller model size** – lower VRAM requirements than the original model +* ⚡ **Compatible with ComfyUI** – drop-in support with custom nodes +* 🔧 **GPU-accelerated inference** – optimized for CUDA 12.1+ + +--- + +## Installation + +### Requirements + +* [ComfyUI](https://github.com/comfyanonymous/ComfyUI) installed +* NVIDIA GPU with **CUDA 12.1+** + +### Setup + +1. Install dependencies: + + ```bash + pip install -r requirements.txt + ``` + +2. Install the DFloat11 custom nodes in ComfyUI: + + ```bash + cd /custom_nodes + git clone https://github.com/LeanModels/ComfyUI-DFloat11.git + ``` + +--- + +## Usage + +1. Once installed, the DFloat11 nodes show up under the `DFloat11` folder in the *Node Library*. +2. Download a DFloat11 model for ComfyUI from [Hugging Face](https://huggingface.co/DFloat11) and place it under `/models/diffusion_models`. +3. Use the `DFloat11 Model Loader` node to load the model in *.safetensors format, which acts as a drop-in replacement for the `Load Diffusion Model` node. + +--- + +## Resources + +* 📖 [DFloat11 Paper (arXiv)](https://arxiv.org/abs/2504.11651) +* 🤗 [DFloat11 Models on Hugging Face](https://huggingface.co/DFloat11) + +--- + +## Contributing + +Contributions are welcome! + +* Open an issue to request new model support +* Submit pull requests for bug fixes or improvements + +[![Contributors](https://contrib.rocks/image?repo=LeanModels/ComfyUI-DFloat11)](https://github.com/LeanModels/ComfyUI-DFloat11/graphs/contributors) diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..4d9f7d8 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .dfloat11_model_loader import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/dfloat11_model_loader.py b/dfloat11_model_loader.py new file mode 100644 index 0000000..b93ad79 --- /dev/null +++ b/dfloat11_model_loader.py @@ -0,0 +1,81 @@ +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", +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..9eecb22 --- /dev/null +++ b/requirements.txt @@ -0,0 +1 @@ +dfloat11[cuda12] \ No newline at end of file