Chroma1-HD-DF11-ComfyUI

Further update: I have uploaded an updated version of this model, that should further reduce disk size and VRAM usage by ~82 MB. This is because I missed out on compressing a small portion of the model (the distilled_guidance_layers) in my original upload. There is <u>no need to download again</u> if you are not having any issues with the older version, but please either update to the latest version of my DF11 custom node fork, or update the patched pattern_dict as shown below if you are still using the official node. I apologize for my earlier oversight, and I will update the other Chroma-based models (just the Flash version for ComfyUI, the diffusers-native models come bundled with the correct pattern_dict in the config file) soon.

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": {
            r"distilled_guidance_layer\.layers\.\d+": (
                "in_layer",
                "out_layer"        
            ),
            r"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",
            ),
            r"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 = {
    r"distilled_guidance_layer\.layers\.\d+": (
        "in_layer",
        "out_layer"        
    ),
    r"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",
    ),
    r"single_blocks\.\d+": (
        "linear1",
        "linear2",
    ),
}