mingyi456/Chroma1-HD-DF11-ComfyUI

Update: I have created a fork of the official DF11 ComfyUI node, that aims to support more models, as well as features: https://github.com/mingyi456/ComfyUI-DFloat11-Extended. It supports the this model out-of-the-box without the need for patching, among other things.

After getting the DF11 compressed versions of Flux to work in ComfyUI, I wanted to try adapting the code to work with other model architectures. Chroma, as expected is easy to do, but I believe others will be harder.

For more information (including how to compress models yourself), check out https://huggingface.co/DFloat11 and https://github.com/LeanModels/DFloat11

Feel free to request for other models for compression as well (for either the diffusers library or ComfyUI), although compressing models that are of architectures that are unfamiliar to me might be more difficult.

How to Use

ComfyUI

Follow the instructions here: https://github.com/LeanModels/ComfyUI-DFloat11. Because the stock DF11 custom node currently only supports Flux-based models, the code needs to be manually patched to support Chroma (I intend to fork the repo and maybe submit a PR at some point, but this shall suffice for now). All that is needed is to explicitly add an entry for "Chroma" in the MODEL_TO_PATTERN_DICT dictionary. The patched code should be as follows:

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",
            ),
        },
        "Chroma": {
            "double_blocks\.\d+": (
                "img_attn.qkv",
                "img_attn.proj",
                "img_mlp.0",
                "img_mlp.2",
                "txt_attn.qkv",
                "txt_attn.proj",
                "txt_mlp.0",
                "txt_mlp.2",
            ),
            "single_blocks\.\d+": (
                "linear1",
                "linear2",
            ),
        },
    }

Finally, use the provided workflow json, or simply replace the "Load Diffusion Model" node of an existing Chroma workflow with the patched "DFloat11 Model Loader" node. If you run into any issues, feel free to leave a comment. The workflow is also embedded in the below png image.

diffusers

Refer to this model instead.

Compression Details

This is the pattern_dict for compressing Chroma-based models in ComfyUI:

pattern_dict_comfyui = {
    "double_blocks\.\d+": (
        "img_attn.qkv",
        "img_attn.proj",
        "img_mlp.0",
        "img_mlp.2",
        "txt_attn.qkv",
        "txt_attn.proj",
        "txt_mlp.0",
        "txt_mlp.2",
    ),
    "single_blocks\.\d+": (
        "linear1",
        "linear2",
    ),
}