mingyi456/SRPO-DF11-ComfyUI

Thanks to @rootey for pointing out that my existing DF11 compressions do not work with ComfyUI. After I realized the issue was with the different naming schemes of the tensors, I worked on trying to translate and remap each tensor in the FluxTransformer2D model (in the diffusers library), to the Flux.diffusion_model model (which ComfyUI, and I believe other platforms, use), only to realize that the tensors corresponding to the DF11 compressed weights expect to be decompressed into a specific structure at compression time, and trying to adapt for it at runtime will be extremely difficult for me, although it should be theoretically possible for an expert to do.

Instead, I figured out a way to manually load an instance of the Flux.diffusion_model model in ComfyUI, and carry out the compression in ComfyUI instead. Unfortunately, this means that 2 separate, mutrally-incompatible versions of DF11 models need to coexist, which I personally find to be bothersome. Hence, I am currently researching on the feasibility of adapting the existing models at runtime to avoid this issue, to avoid having to maintain and upload models in both formats. In the meantime, I plan to release ComfyUI compatible DF11 compressions for only a select handful of models, but feel free to request for any of the existing DF11 models, or new ones, to be remade into ComfyUI format.

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, although compressing models that do not use the Flux architecture might be slightly tricky for me.

This compressed model was made from rockerBOO/flux.1-dev-SRPO's BF16 quantization. Thanks to rockerBOO, without which I would not have been able to directly work with. (My PC only has 48 GB of system RAM, too little to work with a 12B model in FP32 precision)

How to Use

ComfyUI

Follow the instructions here: https://github.com/LeanModels/ComfyUI-DFloat11. After installing the DF11 custom node, use the provided workflow json, or simply replace the "Load Diffusion Model" node of an existing Flux workflow with the "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 Flux-based models in ComfyUI:

pattern_dict_comfyui = {
    "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",
    ),
}